1use serde::{Deserialize, Serialize};
2use std::{collections::BTreeSet, error::Error, fmt};
3
4pub use kcode_speaker_v3_llm_protocol::{
5 GEMINI_FEATURE_PROMPT_ONE, GEMINI_FEATURE_PROMPT_ONE_REVISION, GEMINI_FEATURE_PROMPT_REVISIONS,
6 GEMINI_FEATURE_PROMPT_THREE, GEMINI_FEATURE_PROMPT_THREE_REVISION, GEMINI_FEATURE_PROMPT_TWO,
7 GEMINI_FEATURE_PROMPT_TWO_REVISION, GEMINI_TRANSCRIPT_PROMPT,
8 GEMINI_TRANSCRIPT_PROMPT_REVISION, GPT_STRUCTURING_PROMPT, GPT_STRUCTURING_PROMPT_REVISION,
9};
10pub use kcode_speaker_v3_schema::{
11 FEATURE_NAMES, FEATURE_SCHEMA_REVISION, FeatureVector24, LocalSpeakerLabel,
12 MAX_AUDIO_DURATION_MS, OGG_MEDIA_TYPE, OggAudioMetadata, StructuredAnalysis, StructuredSpeaker,
13 ValidationError, VocalGenderPresentation,
14};
15
16#[cfg(any(feature = "providers", test))]
17use kcode_speaker_v3_llm_protocol::SpeakerFeatureEvidence;
18#[cfg(any(feature = "providers", test))]
19use std::{future::Future, pin::Pin};
20
21const GEMINI_MODEL_ID: &str = "gemini-3.1-pro-preview";
22const TERRA_MODEL_ID: &str = "gpt-5.6-terra";
23
24#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
25pub struct GeminiCohort {
26 pub model_id: String,
27 pub transcript_prompt_revision: String,
28 pub feature_prompt_revisions: [String; 3],
29 pub feature_schema_revision: String,
30}
31
32impl GeminiCohort {
33 pub fn new(model_id: impl Into<String>) -> Self {
34 Self {
35 model_id: model_id.into(),
36 transcript_prompt_revision: GEMINI_TRANSCRIPT_PROMPT_REVISION.into(),
37 feature_prompt_revisions: GEMINI_FEATURE_PROMPT_REVISIONS.map(str::to_owned),
38 feature_schema_revision: FEATURE_SCHEMA_REVISION.into(),
39 }
40 }
41
42 pub fn validate(&self) -> Result<(), ValidationError> {
43 validate_text(&self.model_id, "gemini_model_id")?;
44 validate_text(
45 &self.transcript_prompt_revision,
46 "transcript_prompt_revision",
47 )?;
48 for revision in &self.feature_prompt_revisions {
49 validate_text(revision, "feature_prompt_revision")?;
50 }
51 validate_text(&self.feature_schema_revision, "feature_schema_revision")
52 }
53}
54
55#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
56pub struct StructurerProvenance {
57 pub model_id: String,
58 pub prompt_revision: String,
59}
60
61impl StructurerProvenance {
62 pub fn new(model_id: impl Into<String>) -> Self {
63 Self {
64 model_id: model_id.into(),
65 prompt_revision: GPT_STRUCTURING_PROMPT_REVISION.into(),
66 }
67 }
68
69 pub fn validate(&self) -> Result<(), ValidationError> {
70 validate_text(&self.model_id, "structurer_model_id")?;
71 validate_text(&self.prompt_revision, "structurer_prompt_revision")
72 }
73}
74
75#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
76pub struct AnalysisEnvelope {
77 pub audio: OggAudioMetadata,
78 pub analysis: StructuredAnalysis,
79 pub gemini: GeminiCohort,
80 pub structurer: StructurerProvenance,
81}
82
83impl AnalysisEnvelope {
84 pub fn validate(&self) -> Result<(), ValidationError> {
85 self.audio.validate()?;
86 self.analysis.validate()?;
87 self.gemini.validate()?;
88 self.structurer.validate()
89 }
90}
91
92#[derive(Debug, Clone, PartialEq, Eq)]
93pub enum AnalysisError {
94 Input(String),
95 GeminiTranscript(String),
96 TerraLabels(String),
97 GeminiCache(String),
98 GeminiFeature {
99 speaker: LocalSpeakerLabel,
100 packet: u8,
101 message: String,
102 },
103 TerraStructuring(String),
104 TranscriptMismatch,
105 SpeakerSetMismatch,
106}
107
108impl fmt::Display for AnalysisError {
109 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
110 match self {
111 Self::Input(message) => write!(formatter, "invalid input: {message}"),
112 Self::GeminiTranscript(message) => {
113 write!(formatter, "Gemini transcript failed: {message}")
114 }
115 Self::TerraLabels(message) => {
116 write!(
117 formatter,
118 "Terra speaker-label extraction failed: {message}"
119 )
120 }
121 Self::GeminiCache(message) => {
122 write!(formatter, "Gemini feature cache creation failed: {message}")
123 }
124 Self::GeminiFeature {
125 speaker,
126 packet,
127 message,
128 } => write!(
129 formatter,
130 "Gemini feature call failed for {speaker}, packet {packet}: {message}"
131 ),
132 Self::TerraStructuring(message) => {
133 write!(formatter, "Terra final structuring failed: {message}")
134 }
135 Self::TranscriptMismatch => {
136 formatter.write_str("Terra returned a different transcript")
137 }
138 Self::SpeakerSetMismatch => {
139 formatter.write_str("Terra returned a different speaker set")
140 }
141 }
142 }
143}
144
145impl Error for AnalysisError {}
146
147#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
148pub struct ExecutedAnalysis {
149 pub envelope: AnalysisEnvelope,
150 pub label_extractor: StructurerProvenance,
151}
152
153#[cfg(feature = "providers")]
154pub struct Analyzer {
155 operations: ProviderOperations,
156}
157
158#[cfg(feature = "providers")]
159impl Analyzer {
160 pub fn new(
161 gemini: kcode_gemini_3_1_pro::Gemini31Pro,
162 terra: kcode_codex_terra::CodexTerra,
163 ) -> Self {
164 Self {
165 operations: ProviderOperations {
166 gemini: kcode_speaker_v3_gemini_analysis::GeminiAnalysis::new(gemini),
167 terra: kcode_speaker_v3_terra_analysis::TerraAnalysis::new(terra),
168 },
169 }
170 }
171
172 pub async fn analyze_ogg(
173 &self,
174 bytes: &[u8],
175 duration_ms: u64,
176 filename: Option<String>,
177 ) -> Result<ExecutedAnalysis, AnalysisError> {
178 execute(&self.operations, bytes, duration_ms, filename).await
179 }
180}
181
182fn validate_text(value: &str, field: &'static str) -> Result<(), ValidationError> {
183 (!value.trim().is_empty())
184 .then_some(())
185 .ok_or(ValidationError::Blank(field))
186}
187
188#[cfg(any(feature = "providers", test))]
189type AnalysisFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
190
191#[cfg(any(feature = "providers", test))]
192trait AnalysisOperations: Sync {
193 fn transcript<'a>(
194 &'a self,
195 audio: &'a [u8],
196 ) -> AnalysisFuture<'a, Result<String, AnalysisError>>;
197
198 fn speaker_labels<'a>(
199 &'a self,
200 transcript: &'a str,
201 ) -> AnalysisFuture<'a, Result<Vec<LocalSpeakerLabel>, AnalysisError>>;
202
203 fn feature_evidence<'a>(
204 &'a self,
205 audio: &'a [u8],
206 transcript: &'a str,
207 labels: &'a [LocalSpeakerLabel],
208 ) -> AnalysisFuture<'a, Result<Vec<SpeakerFeatureEvidence>, AnalysisError>>;
209
210 fn structured_analysis<'a>(
211 &'a self,
212 transcript: &'a str,
213 evidence: Vec<SpeakerFeatureEvidence>,
214 ) -> AnalysisFuture<'a, Result<StructuredAnalysis, AnalysisError>>;
215}
216
217#[cfg(any(feature = "providers", test))]
218struct ProviderOperations {
219 gemini: kcode_speaker_v3_gemini_analysis::GeminiAnalysis,
220 terra: kcode_speaker_v3_terra_analysis::TerraAnalysis,
221}
222
223#[cfg(any(feature = "providers", test))]
224impl AnalysisOperations for ProviderOperations {
225 fn transcript<'a>(
226 &'a self,
227 audio: &'a [u8],
228 ) -> AnalysisFuture<'a, Result<String, AnalysisError>> {
229 Box::pin(async move {
230 self.gemini.transcript(audio).await.map_err(|error| match error {
231 kcode_speaker_v3_gemini_analysis::GeminiTranscriptError::Provider(message)
232 | kcode_speaker_v3_gemini_analysis::GeminiTranscriptError::Protocol(message) => {
233 AnalysisError::GeminiTranscript(message)
234 }
235 })
236 })
237 }
238
239 fn speaker_labels<'a>(
240 &'a self,
241 transcript: &'a str,
242 ) -> AnalysisFuture<'a, Result<Vec<LocalSpeakerLabel>, AnalysisError>> {
243 Box::pin(async move {
244 self.terra
245 .speaker_labels(transcript)
246 .await
247 .map_err(|error| match error {
248 kcode_speaker_v3_terra_analysis::TerraAnalysisError::Protocol(message)
249 | kcode_speaker_v3_terra_analysis::TerraAnalysisError::Provider(message) => {
250 AnalysisError::TerraLabels(message)
251 }
252 })
253 })
254 }
255
256 fn feature_evidence<'a>(
257 &'a self,
258 audio: &'a [u8],
259 transcript: &'a str,
260 labels: &'a [LocalSpeakerLabel],
261 ) -> AnalysisFuture<'a, Result<Vec<SpeakerFeatureEvidence>, AnalysisError>> {
262 Box::pin(async move {
263 self.gemini
264 .feature_evidence(audio, transcript, labels)
265 .await
266 .map_err(|error| match error {
267 kcode_speaker_v3_gemini_analysis::GeminiFeatureError::Cache(message) => {
268 AnalysisError::GeminiCache(message)
269 }
270 kcode_speaker_v3_gemini_analysis::GeminiFeatureError::Feature {
271 speaker,
272 packet,
273 message,
274 } => AnalysisError::GeminiFeature {
275 speaker,
276 packet,
277 message,
278 },
279 })
280 })
281 }
282
283 fn structured_analysis<'a>(
284 &'a self,
285 transcript: &'a str,
286 evidence: Vec<SpeakerFeatureEvidence>,
287 ) -> AnalysisFuture<'a, Result<StructuredAnalysis, AnalysisError>> {
288 Box::pin(async move {
289 self.terra
290 .structured_analysis(transcript, evidence)
291 .await
292 .map_err(|error| match error {
293 kcode_speaker_v3_terra_analysis::TerraAnalysisError::Protocol(message)
294 | kcode_speaker_v3_terra_analysis::TerraAnalysisError::Provider(message) => {
295 AnalysisError::TerraStructuring(message)
296 }
297 })
298 })
299 }
300}
301
302#[cfg(any(feature = "providers", test))]
303async fn execute<O: AnalysisOperations>(
304 operations: &O,
305 bytes: &[u8],
306 duration_ms: u64,
307 filename: Option<String>,
308) -> Result<ExecutedAnalysis, AnalysisError> {
309 let audio = OggAudioMetadata::from_bytes(bytes, duration_ms, filename)
310 .map_err(|error| AnalysisError::Input(error.to_string()))?;
311 let transcript = operations.transcript(bytes).await?;
312 let labels = operations.speaker_labels(&transcript).await?;
313 let evidence = operations
314 .feature_evidence(bytes, &transcript, &labels)
315 .await?;
316 let analysis = operations
317 .structured_analysis(&transcript, evidence)
318 .await?;
319
320 if analysis.transcript != transcript {
321 return Err(AnalysisError::TranscriptMismatch);
322 }
323
324 let expected_speakers = labels.iter().copied().collect::<BTreeSet<_>>();
325 let returned_speakers = analysis
326 .speakers
327 .iter()
328 .map(|speaker| speaker.speaker)
329 .collect::<BTreeSet<_>>();
330 if expected_speakers != returned_speakers {
331 return Err(AnalysisError::SpeakerSetMismatch);
332 }
333
334 let envelope = AnalysisEnvelope {
335 audio,
336 analysis,
337 gemini: GeminiCohort::new(GEMINI_MODEL_ID),
338 structurer: StructurerProvenance::new(TERRA_MODEL_ID),
339 };
340 envelope
341 .validate()
342 .map_err(|error| AnalysisError::TerraStructuring(error.to_string()))?;
343
344 let label_extractor = StructurerProvenance {
345 model_id: TERRA_MODEL_ID.into(),
346 prompt_revision: kcode_speaker_v3_llm_protocol::TERRA_SPEAKER_LABELS_PROMPT_REVISION.into(),
347 };
348 label_extractor
349 .validate()
350 .map_err(|error| AnalysisError::TerraLabels(error.to_string()))?;
351
352 Ok(ExecutedAnalysis {
353 envelope,
354 label_extractor,
355 })
356}
357
358#[cfg(test)]
359mod tests {
360 use super::*;
361 use futures::{executor::block_on, future::poll_fn, join};
362 use std::{
363 sync::{
364 Arc, Mutex,
365 atomic::{AtomicBool, AtomicUsize, Ordering},
366 },
367 task::Poll,
368 time::Instant,
369 };
370
371 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
372 enum FailureStage {
373 Transcript,
374 Labels,
375 Features,
376 Final,
377 }
378
379 #[derive(Clone)]
380 struct FakeConfig {
381 transcript: String,
382 labels: Vec<LocalSpeakerLabel>,
383 analysis: StructuredAnalysis,
384 failure: Option<FailureStage>,
385 wait_for: Option<Arc<AtomicBool>>,
386 mark_complete: Option<Arc<AtomicBool>>,
387 }
388
389 struct FakeState {
390 config: FakeConfig,
391 transcript_calls: AtomicUsize,
392 label_calls: AtomicUsize,
393 feature_calls: AtomicUsize,
394 final_calls: AtomicUsize,
395 calls: Mutex<Vec<&'static str>>,
396 }
397
398 #[derive(Clone)]
399 struct FakeOperations {
400 state: Arc<FakeState>,
401 }
402
403 impl FakeOperations {
404 fn successful(speaker_count: u32) -> Self {
405 let transcript = "[high] Speaker 1: exact transcript".to_owned();
406 Self::from_config(FakeConfig {
407 labels: (1..=speaker_count).map(label).collect(),
408 analysis: structured_analysis(&transcript, speaker_count),
409 transcript,
410 failure: None,
411 wait_for: None,
412 mark_complete: None,
413 })
414 }
415
416 fn from_config(config: FakeConfig) -> Self {
417 Self {
418 state: Arc::new(FakeState {
419 config,
420 transcript_calls: AtomicUsize::new(0),
421 label_calls: AtomicUsize::new(0),
422 feature_calls: AtomicUsize::new(0),
423 final_calls: AtomicUsize::new(0),
424 calls: Mutex::new(Vec::new()),
425 }),
426 }
427 }
428
429 fn with_config(&self, update: impl FnOnce(&mut FakeConfig)) -> Self {
430 let mut config = self.state.config.clone();
431 update(&mut config);
432 Self::from_config(config)
433 }
434 }
435
436 impl AnalysisOperations for FakeOperations {
437 fn transcript<'a>(
438 &'a self,
439 _audio: &'a [u8],
440 ) -> AnalysisFuture<'a, Result<String, AnalysisError>> {
441 Box::pin(async move {
442 self.state.transcript_calls.fetch_add(1, Ordering::SeqCst);
443 self.state.calls.lock().unwrap().push("transcript");
444 if let Some(wait_for) = &self.state.config.wait_for {
445 poll_fn(|context| {
446 if wait_for.load(Ordering::SeqCst) {
447 Poll::Ready(())
448 } else {
449 context.waker().wake_by_ref();
450 Poll::Pending
451 }
452 })
453 .await;
454 }
455 if self.state.config.failure == Some(FailureStage::Transcript) {
456 return Err(AnalysisError::GeminiTranscript("transcript".into()));
457 }
458 Ok(self.state.config.transcript.clone())
459 })
460 }
461
462 fn speaker_labels<'a>(
463 &'a self,
464 _transcript: &'a str,
465 ) -> AnalysisFuture<'a, Result<Vec<LocalSpeakerLabel>, AnalysisError>> {
466 Box::pin(async move {
467 self.state.label_calls.fetch_add(1, Ordering::SeqCst);
468 self.state.calls.lock().unwrap().push("labels");
469 if self.state.config.failure == Some(FailureStage::Labels) {
470 return Err(AnalysisError::TerraLabels("labels".into()));
471 }
472 Ok(self.state.config.labels.clone())
473 })
474 }
475
476 fn feature_evidence<'a>(
477 &'a self,
478 _audio: &'a [u8],
479 _transcript: &'a str,
480 labels: &'a [LocalSpeakerLabel],
481 ) -> AnalysisFuture<'a, Result<Vec<SpeakerFeatureEvidence>, AnalysisError>> {
482 Box::pin(async move {
483 self.state.feature_calls.fetch_add(1, Ordering::SeqCst);
484 self.state.calls.lock().unwrap().push("features");
485 if self.state.config.failure == Some(FailureStage::Features) {
486 return Err(AnalysisError::GeminiFeature {
487 speaker: label(1),
488 packet: 2,
489 message: "features".into(),
490 });
491 }
492 labels
493 .iter()
494 .copied()
495 .map(|speaker| {
496 SpeakerFeatureEvidence::new(
497 speaker,
498 format!("{speaker} packet 1"),
499 format!("{speaker} packet 2"),
500 format!("{speaker} packet 3"),
501 )
502 .map_err(|error| AnalysisError::GeminiFeature {
503 speaker,
504 packet: 1,
505 message: error.to_string(),
506 })
507 })
508 .collect()
509 })
510 }
511
512 fn structured_analysis<'a>(
513 &'a self,
514 _transcript: &'a str,
515 _evidence: Vec<SpeakerFeatureEvidence>,
516 ) -> AnalysisFuture<'a, Result<StructuredAnalysis, AnalysisError>> {
517 Box::pin(async move {
518 self.state.final_calls.fetch_add(1, Ordering::SeqCst);
519 self.state.calls.lock().unwrap().push("final");
520 if self.state.config.failure == Some(FailureStage::Final) {
521 return Err(AnalysisError::TerraStructuring("final".into()));
522 }
523 if let Some(mark_complete) = &self.state.config.mark_complete {
524 mark_complete.store(true, Ordering::SeqCst);
525 }
526 Ok(self.state.config.analysis.clone())
527 })
528 }
529 }
530
531 fn label(number: u32) -> LocalSpeakerLabel {
532 LocalSpeakerLabel::new(number).unwrap()
533 }
534
535 fn structured_analysis(transcript: &str, speaker_count: u32) -> StructuredAnalysis {
536 StructuredAnalysis {
537 transcript: transcript.into(),
538 speakers: (1..=speaker_count)
539 .map(|number| StructuredSpeaker {
540 speaker: label(number),
541 language: "English".into(),
542 features: FeatureVector24::default(),
543 features_usable_for_training: false,
544 })
545 .collect(),
546 }
547 }
548
549 fn ogg() -> Vec<u8> {
550 let mut bytes = vec![0; 28];
551 bytes[..4].copy_from_slice(b"OggS");
552 bytes[4] = 0;
553 bytes[26] = 1;
554 bytes[27] = 0;
555 bytes
556 }
557
558 #[test]
559 fn zero_one_and_many_speaker_workflows_keep_exact_stage_order() {
560 for speaker_count in [0, 1, 40] {
561 let operations = FakeOperations::successful(speaker_count);
562 let result =
563 block_on(execute(&operations, &ogg(), 1, Some("voice.ogg".into()))).unwrap();
564 assert_eq!(
565 result.envelope.analysis.speakers.len(),
566 speaker_count as usize
567 );
568 assert_eq!(
569 *operations.state.calls.lock().unwrap(),
570 ["transcript", "labels", "features", "final"]
571 );
572 assert_eq!(
573 result.label_extractor,
574 StructurerProvenance {
575 model_id: TERRA_MODEL_ID.into(),
576 prompt_revision:
577 kcode_speaker_v3_llm_protocol::TERRA_SPEAKER_LABELS_PROMPT_REVISION.into(),
578 }
579 );
580 }
581 }
582
583 #[test]
584 fn input_and_each_provider_stage_fail_without_retry() {
585 let input = FakeOperations::successful(1);
586 assert!(matches!(
587 block_on(execute(&input, b"bad", 1, None)),
588 Err(AnalysisError::Input(_))
589 ));
590 assert!(input.state.calls.lock().unwrap().is_empty());
591
592 for (stage, expected) in [
593 (FailureStage::Transcript, vec!["transcript"]),
594 (FailureStage::Labels, vec!["transcript", "labels"]),
595 (
596 FailureStage::Features,
597 vec!["transcript", "labels", "features"],
598 ),
599 (
600 FailureStage::Final,
601 vec!["transcript", "labels", "features", "final"],
602 ),
603 ] {
604 let operations =
605 FakeOperations::successful(1).with_config(|config| config.failure = Some(stage));
606 assert!(block_on(execute(&operations, &ogg(), 1, None)).is_err());
607 assert_eq!(*operations.state.calls.lock().unwrap(), expected);
608 }
609 }
610
611 #[test]
612 fn cross_stage_transcript_and_speaker_mismatches_are_rejected() {
613 let transcript = FakeOperations::successful(1).with_config(|config| {
614 config.analysis = structured_analysis("different", 1);
615 });
616 assert_eq!(
617 block_on(execute(&transcript, &ogg(), 1, None)),
618 Err(AnalysisError::TranscriptMismatch)
619 );
620
621 let speakers = FakeOperations::successful(1).with_config(|config| {
622 config.analysis = structured_analysis(&config.transcript, 2);
623 });
624 assert_eq!(
625 block_on(execute(&speakers, &ogg(), 1, None)),
626 Err(AnalysisError::SpeakerSetMismatch)
627 );
628 }
629
630 #[test]
631 fn a_blocked_analysis_does_not_block_an_unrelated_analysis() {
632 let completed = Arc::new(AtomicBool::new(false));
633 let fast = FakeOperations::successful(0).with_config(|config| {
634 config.mark_complete = Some(completed.clone());
635 });
636 let slow = FakeOperations::successful(0).with_config(|config| {
637 config.wait_for = Some(completed.clone());
638 });
639 let slow_audio = ogg();
640 let fast_audio = ogg();
641 let (slow_result, fast_result) = block_on(async {
642 join!(
643 execute(&slow, &slow_audio, 1, None),
644 execute(&fast, &fast_audio, 1, None)
645 )
646 });
647 slow_result.unwrap();
648 fast_result.unwrap();
649 assert!(completed.load(Ordering::SeqCst));
650 }
651
652 #[test]
653 fn provenance_preserves_the_previous_public_contract() {
654 let cohort = GeminiCohort::new("gemini-model");
655 assert_eq!(
656 cohort.feature_prompt_revisions,
657 GEMINI_FEATURE_PROMPT_REVISIONS.map(str::to_owned)
658 );
659 cohort.validate().unwrap();
660 StructurerProvenance::new("gpt-5.6").validate().unwrap();
661 assert_eq!(
662 GeminiCohort::new(" ").validate(),
663 Err(ValidationError::Blank("gemini_model_id"))
664 );
665 }
666
667 #[test]
668 fn reference_scale_local_orchestration_completes_within_envelope() {
669 let started = Instant::now();
670 let operations = FakeOperations::successful(1000);
671 let result = block_on(execute(&operations, &ogg(), 1, None)).unwrap();
672 assert_eq!(result.envelope.analysis.speakers.len(), 1000);
673 assert!(started.elapsed().as_secs() < 10);
674 }
675
676 #[test]
677 fn concrete_provider_operations_compile() {
678 fn require_operations<O: AnalysisOperations>() {}
679 require_operations::<ProviderOperations>();
680 }
681}