1use eredu_core::{RealtimeFrameForcing, RealtimeInputFrame, RealtimeSpeechConfig};
4
5use crate::TokenDomain;
6
7#[derive(Debug, Clone, Eq, PartialEq)]
9pub struct RealtimeIngressContract {
10 schedule: RealtimeSpeechConfig,
11 text: TokenDomain,
12 audio: TokenDomain,
13}
14
15impl RealtimeIngressContract {
16 pub fn new(
18 schedule: RealtimeSpeechConfig,
19 text: TokenDomain,
20 audio: TokenDomain,
21 ) -> Result<Self, RealtimeIngressError> {
22 validate_token(schedule.text_padding_token(), text, RealtimeTokenKind::Text)?;
23 validate_token(
24 schedule.audio_padding_token(),
25 audio,
26 RealtimeTokenKind::Audio,
27 )?;
28 Ok(Self {
29 schedule,
30 text,
31 audio,
32 })
33 }
34
35 pub const fn schedule(&self) -> &RealtimeSpeechConfig {
37 &self.schedule
38 }
39
40 pub const fn text_domain(&self) -> TokenDomain {
42 self.text
43 }
44
45 pub const fn audio_domain(&self) -> TokenDomain {
47 self.audio
48 }
49
50 pub fn validate<'a>(
52 &'a self,
53 frame: &'a RealtimeInputFrame,
54 ) -> Result<ValidatedRealtimeInput<'a>, RealtimeIngressError> {
55 let batch = frame.batch();
56 if batch == 0 {
57 return Err(RealtimeIngressError::EmptyBatch);
58 }
59 let input_columns = self.schedule.input_audio_codebooks();
60 validate_shape(
61 RealtimePayloadKind::InputAudio,
62 frame.input_audio_tokens().len(),
63 batch,
64 input_columns,
65 )?;
66 for &token in frame.input_audio_tokens() {
67 validate_token(token, self.audio, RealtimeTokenKind::Audio)?;
68 }
69
70 let generated = self.schedule.generated_audio_codebooks();
71 let forced_audio = frame.forced_generated_audio_tokens();
72 let forcing = match (forced_audio, frame.forced_generated_audio_codebooks()) {
73 (None, None) => vec![false; generated],
74 (None, Some(_)) => return Err(RealtimeIngressError::ForcingMaskWithoutPayload),
75 (Some(tokens), mask) => {
76 validate_shape(
77 RealtimePayloadKind::ForcedAudio,
78 tokens.len(),
79 batch,
80 generated,
81 )?;
82 let mask = mask.map_or_else(|| vec![true; generated], <[bool]>::to_vec);
83 if mask.len() != generated {
84 return Err(RealtimeIngressError::ForcingMaskCount {
85 expected: generated,
86 actual: mask.len(),
87 });
88 }
89 for row in tokens.chunks_exact(generated) {
90 for (token, selected) in row.iter().zip(&mask) {
91 if *selected {
92 validate_token(*token, self.audio, RealtimeTokenKind::Audio)?;
93 }
94 }
95 }
96 mask
97 }
98 };
99
100 let forced_text = frame.forced_text_tokens();
101 if let Some(tokens) = forced_text {
102 validate_shape(RealtimePayloadKind::ForcedText, tokens.len(), batch, 1)?;
103 for &token in tokens {
104 validate_token(token, self.text, RealtimeTokenKind::Text)?;
105 }
106 }
107 Ok(ValidatedRealtimeInput {
108 contract: self,
109 frame,
110 forcing: RealtimeFrameForcing::new(forced_text.is_some(), forcing),
111 })
112 }
113}
114
115pub struct ValidatedRealtimeInput<'a> {
117 contract: &'a RealtimeIngressContract,
118 frame: &'a RealtimeInputFrame,
119 forcing: RealtimeFrameForcing,
120}
121
122impl ValidatedRealtimeInput<'_> {
123 pub const fn frame(&self) -> &RealtimeInputFrame {
125 self.frame
126 }
127
128 pub const fn forcing(&self) -> &RealtimeFrameForcing {
130 &self.forcing
131 }
132
133 pub fn materialize<M: RealtimeHostTokenMaterializer>(
135 &self,
136 materializer: &mut M,
137 ) -> Result<MaterializedRealtimeInput<M::Tensor>, M::Error> {
138 let batch = self.frame.batch();
139 let input_audio = materializer.materialize_i32(
140 self.frame.input_audio_tokens(),
141 [batch, self.contract.schedule.input_audio_codebooks()],
142 )?;
143 let forced_audio = self
144 .frame
145 .forced_generated_audio_tokens()
146 .map(|tokens| {
147 materializer.materialize_i32(
148 tokens,
149 [batch, self.contract.schedule.generated_audio_codebooks()],
150 )
151 })
152 .transpose()?;
153 let forced_text = self
154 .frame
155 .forced_text_tokens()
156 .map(|tokens| materializer.materialize_i32(tokens, [batch, 1]))
157 .transpose()?;
158 Ok(MaterializedRealtimeInput {
159 schedule: self.contract.schedule.clone(),
160 batch,
161 input_audio,
162 forced_audio,
163 forced_text,
164 forcing: self.forcing.clone(),
165 retain_diagnostics: self.frame.retains_diagnostics(),
166 })
167 }
168}
169
170pub trait RealtimeHostTokenMaterializer {
172 type Tensor;
174 type Error;
176
177 fn materialize_i32(
179 &mut self,
180 values: &[i32],
181 shape: [usize; 2],
182 ) -> Result<Self::Tensor, Self::Error>;
183}
184
185pub struct MaterializedRealtimeInput<T> {
187 schedule: RealtimeSpeechConfig,
188 batch: usize,
189 input_audio: T,
190 forced_audio: Option<T>,
191 forced_text: Option<T>,
192 forcing: RealtimeFrameForcing,
193 retain_diagnostics: bool,
194}
195
196impl<T> MaterializedRealtimeInput<T> {
197 pub const fn schedule(&self) -> &RealtimeSpeechConfig {
199 &self.schedule
200 }
201
202 pub const fn batch(&self) -> usize {
204 self.batch
205 }
206
207 pub const fn input_audio(&self) -> &T {
209 &self.input_audio
210 }
211
212 pub const fn forced_audio(&self) -> Option<&T> {
214 self.forced_audio.as_ref()
215 }
216
217 pub const fn forced_text(&self) -> Option<&T> {
219 self.forced_text.as_ref()
220 }
221
222 pub const fn forcing(&self) -> &RealtimeFrameForcing {
224 &self.forcing
225 }
226
227 pub const fn retains_diagnostics(&self) -> bool {
229 self.retain_diagnostics
230 }
231
232 pub fn into_parts(
234 self,
235 ) -> (
236 RealtimeSpeechConfig,
237 usize,
238 T,
239 Option<T>,
240 Option<T>,
241 RealtimeFrameForcing,
242 bool,
243 ) {
244 (
245 self.schedule,
246 self.batch,
247 self.input_audio,
248 self.forced_audio,
249 self.forced_text,
250 self.forcing,
251 self.retain_diagnostics,
252 )
253 }
254}
255
256#[derive(Debug, Clone, Copy, Eq, PartialEq)]
257pub enum RealtimePayloadKind {
259 InputAudio,
261 ForcedAudio,
263 ForcedText,
265}
266
267#[derive(Debug, Clone, Copy, Eq, PartialEq)]
268pub enum RealtimeTokenKind {
270 Text,
272 Audio,
274}
275
276fn validate_shape(
277 payload: RealtimePayloadKind,
278 actual: usize,
279 rows: usize,
280 columns: usize,
281) -> Result<(), RealtimeIngressError> {
282 let expected = rows
283 .checked_mul(columns)
284 .ok_or(RealtimeIngressError::ShapeOverflow { rows, columns })?;
285 if actual == expected {
286 Ok(())
287 } else {
288 Err(RealtimeIngressError::PayloadShape {
289 payload,
290 expected,
291 actual,
292 })
293 }
294}
295
296fn validate_token(
297 token: i32,
298 domain: TokenDomain,
299 kind: RealtimeTokenKind,
300) -> Result<(), RealtimeIngressError> {
301 let token = usize::try_from(token).map_err(|_| RealtimeIngressError::TokenDomain {
302 kind,
303 token,
304 cardinality: domain.cardinality(),
305 })?;
306 if token < domain.cardinality() {
307 Ok(())
308 } else {
309 Err(RealtimeIngressError::TokenDomain {
310 kind,
311 token: i32::try_from(token).unwrap_or(i32::MAX),
312 cardinality: domain.cardinality(),
313 })
314 }
315}
316
317#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
319#[non_exhaustive]
320pub enum RealtimeIngressError {
321 #[error("realtime input batch must be positive")]
323 EmptyBatch,
324 #[error("realtime payload shape {rows}x{columns} overflowed")]
326 ShapeOverflow {
327 rows: usize,
329 columns: usize,
331 },
332 #[error("realtime {payload:?} payload has {actual} values, expected {expected}")]
334 PayloadShape {
335 payload: RealtimePayloadKind,
337 expected: usize,
339 actual: usize,
341 },
342 #[error("realtime generated-audio forcing mask has no payload")]
344 ForcingMaskWithoutPayload,
345 #[error("realtime forcing mask has {actual} entries, expected {expected}")]
347 ForcingMaskCount {
348 expected: usize,
350 actual: usize,
352 },
353 #[error("realtime {kind:?} token {token} is outside 0..{cardinality}")]
355 TokenDomain {
356 kind: RealtimeTokenKind,
358 token: i32,
360 cardinality: usize,
362 },
363}
364
365#[cfg(test)]
366mod tests {
367 use std::convert::Infallible;
368
369 use eredu_core::{RealtimeFrameConvention, RealtimeInputFrame};
370
371 use super::*;
372
373 fn contract() -> RealtimeIngressContract {
374 RealtimeIngressContract::new(
375 RealtimeSpeechConfig::new(
376 4,
377 2,
378 2,
379 2,
380 16,
381 15,
382 RealtimeFrameConvention::FeedbackAlignedHistory,
383 vec![0, 1, 2, 1, 2],
384 )
385 .unwrap(),
386 TokenDomain::new(17),
387 TokenDomain::new(16),
388 )
389 .unwrap()
390 }
391
392 #[derive(Default)]
393 struct Recorder(Vec<(Vec<i32>, [usize; 2])>);
394
395 impl RealtimeHostTokenMaterializer for Recorder {
396 type Tensor = usize;
397 type Error = Infallible;
398
399 fn materialize_i32(
400 &mut self,
401 values: &[i32],
402 shape: [usize; 2],
403 ) -> Result<Self::Tensor, Self::Error> {
404 self.0.push((values.to_vec(), shape));
405 Ok(self.0.len() - 1)
406 }
407 }
408
409 #[test]
410 fn validation_precedes_every_opaque_materialization() {
411 let contract = contract();
412 let invalid = RealtimeInputFrame::new(2, vec![1, 2, 3, 99]);
413 let mut materializer = Recorder::default();
414 assert!(matches!(
415 contract.validate(&invalid),
416 Err(RealtimeIngressError::TokenDomain { .. })
417 ));
418 assert!(materializer.0.is_empty());
419
420 let valid = RealtimeInputFrame::new(2, vec![1, 2, 3, 4])
421 .with_partially_forced_generated_audio(vec![5, 99, 6, 99], vec![true, false])
422 .with_forced_text(vec![7, 8])
423 .with_diagnostics();
424 let validated = contract.validate(&valid).unwrap();
425 assert_eq!(validated.forcing().generated_audio(), &[true, false]);
426 let input = validated.materialize(&mut materializer).unwrap();
427 assert_eq!(materializer.0.len(), 3);
428 assert_eq!(materializer.0[0].1, [2, 2]);
429 assert_eq!(materializer.0[1].1, [2, 2]);
430 assert_eq!(materializer.0[2].1, [2, 1]);
431 assert!(input.retains_diagnostics());
432 }
433
434 #[test]
435 fn shapes_masks_and_selected_domains_fail_closed() {
436 let contract = contract();
437 assert_eq!(
438 contract.validate(&RealtimeInputFrame::new(0, vec![])).err(),
439 Some(RealtimeIngressError::EmptyBatch)
440 );
441 assert!(matches!(
442 contract.validate(&RealtimeInputFrame::new(2, vec![1, 2, 3])),
443 Err(RealtimeIngressError::PayloadShape { .. })
444 ));
445 assert!(matches!(
446 contract.validate(
447 &RealtimeInputFrame::new(1, vec![1, 2])
448 .with_partially_forced_generated_audio(vec![3, 4], vec![true])
449 ),
450 Err(RealtimeIngressError::ForcingMaskCount { .. })
451 ));
452 assert!(matches!(
453 contract.validate(&RealtimeInputFrame::new(1, vec![1, 2]).with_forced_text(vec![17])),
454 Err(RealtimeIngressError::TokenDomain { .. })
455 ));
456 }
457}