Skip to main content

synapto_interface/
speech_to_text.rs

1use crate::cognitive::CognitiveOutputSpeech;
2use crate::cognitive_output_audio::CognitiveOutputAudio;
3use crate::peer_input_audio::{PEER_INPUT_AUDIO_CHUNK_DURATION, PeerInputAudio};
4use crate::plugin::Plugin;
5use schemars::JsonSchema;
6use serde::{Deserialize, Serialize};
7use std::sync::Arc;
8
9/// Calculates the discrete audio chunk indices for a given continuous time range.
10///
11/// This method converts real-world time (in seconds) into discrete audio chunk sequence numbers.
12/// To prevent dropping short words and correctly handle boundary overlaps, it deliberately applies:
13/// - `floor` to the start time: ensuring the entire chunk where the word begins is included.
14/// - `ceil` to the end time: ensuring the chunk where the word ends is captured, forming a
15///   strict mathematically exclusive boundary `[start_index, end_index)`.
16pub fn calculate_chunk_indices(base_index: u64, start_secs: f64, end_secs: f64) -> (u64, u64) {
17    let chunk_duration_secs = PEER_INPUT_AUDIO_CHUNK_DURATION.as_secs_f64();
18    let start_index = base_index + (start_secs / chunk_duration_secs).floor() as u64;
19    let end_index = base_index + (end_secs / chunk_duration_secs).ceil() as u64;
20    (start_index, end_index)
21}
22
23/// Signal indicating that speech activity has been detected.
24#[derive(Clone)]
25pub struct SpeechDetected(Arc<tokio::sync::Notify>);
26
27impl SpeechDetected {
28    pub fn new(notify: Arc<tokio::sync::Notify>) -> Self {
29        Self(notify)
30    }
31    pub fn notify(&self) {
32        self.0.notify_waiters();
33    }
34}
35
36/// A unique identifier for a speaker.
37#[derive(
38    Serialize,
39    Deserialize,
40    PartialEq,
41    Eq,
42    Hash,
43    Debug,
44    Clone,
45    JsonSchema,
46    derive_more::Display,
47    derive_more::From,
48    derive_more::Deref,
49)]
50pub struct SpeakerId(pub String);
51
52impl SpeakerId {
53    pub fn new(speaker_id: String) -> Self {
54        Self(speaker_id)
55    }
56}
57
58/// Represents a single word within a transcription.
59#[derive(Deserialize, Serialize, Default, Clone, Debug, JsonSchema)]
60pub struct Word {
61    /// The start index of the audio chunk where this word begins.
62    pub start_index: Option<u64>,
63    /// The end index of the audio chunk where this word ends.
64    pub end_index: Option<u64>,
65    /// The text of the word.
66    pub word: String,
67    /// Optional hint about the speaker identity for this specific word.
68    pub speaker_hint: Option<String>,
69}
70
71/// A transcribed segment of speech.
72#[derive(Deserialize, Serialize, Default, Clone, Debug, JsonSchema)]
73pub struct SpeechTranscript {
74    /// The sequence number of the starting audio chunk.
75    pub start_index: u64,
76    /// The sequence number of the ending audio chunk.
77    pub end_index: u64,
78    /// The complete transcribed text for this segment.
79    pub transcript: String,
80    /// Individual words with their respective timing and metadata.
81    pub words: Vec<Word>,
82}
83
84/// Represents an audio chunk with its associated voice activity status.
85#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema)]
86pub enum InputVoiceAudio {
87    /// The chunk contains active voice.
88    Voice(PeerInputAudioIndexed),
89    /// The chunk contains silence or non-voice background noise.
90    NoVoice(PeerInputAudioIndexed),
91}
92
93/// A raw audio chunk bundled with its sequence sequence index.
94#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema)]
95pub struct PeerInputAudioIndexed {
96    /// The raw audio data.
97    pub audio: PeerInputAudio,
98    /// The monotonically increasing sequence number of this chunk.
99    pub index: u64,
100}
101
102#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
103pub struct CosineSimilarity(pub f32);
104
105#[derive(Clone, Debug)]
106pub struct WordOverlap {
107    pub start_index: u64,
108    pub end_index: u64,
109    pub overlaps: std::collections::HashMap<InternalSpeaker, u64>,
110    pub word: String,
111}
112
113pub type SpeakerHeuristicFn =
114    dyn Fn(&[WordOverlap], &[SpeakerSegment]) -> Vec<Option<SpeakerId>> + Send + Sync;
115
116#[derive(Clone)]
117pub struct SpeakerHeuristicCallback(std::sync::Arc<SpeakerHeuristicFn>);
118
119impl std::fmt::Debug for SpeakerHeuristicCallback {
120    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
121        f.debug_struct("SpeakerHeuristicCallback")
122            .finish_non_exhaustive()
123    }
124}
125
126impl SpeakerHeuristicCallback {
127    pub fn new<F>(callback: F) -> Self
128    where
129        F: Fn(&[WordOverlap], &[SpeakerSegment]) -> Vec<Option<SpeakerId>> + Send + Sync + 'static,
130    {
131        Self(std::sync::Arc::new(callback))
132    }
133
134    pub fn evaluate(
135        &self,
136        words: &[WordOverlap],
137        segments: &[SpeakerSegment],
138    ) -> Vec<Option<SpeakerId>> {
139        (self.0)(words, segments)
140    }
141}
142
143#[derive(Serialize, Deserialize, Debug, Clone)]
144pub enum InternalSpeaker {
145    Unknown(Option<(SpeakerId, CosineSimilarity)>),
146    Recognized(SpeakerId),
147}
148
149impl PartialEq for InternalSpeaker {
150    fn eq(&self, other: &Self) -> bool {
151        match (self, other) {
152            (InternalSpeaker::Unknown(a), InternalSpeaker::Unknown(b)) => match (a, b) {
153                (Some((id_a, score_a)), Some((id_b, score_b))) => {
154                    id_a == id_b && score_a.0 == score_b.0
155                }
156                (None, None) => true,
157                _ => false,
158            },
159            (InternalSpeaker::Recognized(a), InternalSpeaker::Recognized(b)) => a == b,
160            _ => false,
161        }
162    }
163}
164
165impl Eq for InternalSpeaker {}
166
167impl std::hash::Hash for InternalSpeaker {
168    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
169        match self {
170            InternalSpeaker::Unknown(opt) => {
171                0_u8.hash(state);
172                if let Some((id, score)) = opt {
173                    1_u8.hash(state);
174                    id.hash(state);
175                    score.0.to_bits().hash(state);
176                } else {
177                    0_u8.hash(state);
178                }
179            }
180            InternalSpeaker::Recognized(id) => {
181                1_u8.hash(state);
182                id.hash(state);
183            }
184        }
185    }
186}
187
188impl From<InternalSpeaker> for crate::peer_input::Speaker {
189    fn from(val: InternalSpeaker) -> Self {
190        match val {
191            InternalSpeaker::Unknown(None) => crate::peer_input::Speaker::Unknown(None),
192            InternalSpeaker::Unknown(Some((id, _))) => crate::peer_input::Speaker::Unknown(Some(
193                SpeakerId(format!("Maybe {}", id).to_string()),
194            )),
195            InternalSpeaker::Recognized(id) => {
196                crate::peer_input::Speaker::Recognized(SpeakerId(id.to_string()))
197            }
198        }
199    }
200}
201
202#[derive(Clone, Debug, Serialize, Deserialize)]
203pub struct SpeakerSegment {
204    pub speaker: InternalSpeaker,
205    pub start_index: u64,
206    pub end_index: u64,
207}
208
209impl From<InputVoiceAudio> for PeerInputAudio {
210    fn from(input_voice_audio: InputVoiceAudio) -> Self {
211        match input_voice_audio {
212            InputVoiceAudio::Voice(peer_input_audio) => peer_input_audio.audio,
213            InputVoiceAudio::NoVoice(peer_input_audio) => peer_input_audio.audio,
214        }
215    }
216}
217
218impl From<InputVoiceAudio> for [i32; crate::peer_input_audio::PEER_INPUT_AUDIO_CHUNK_SIZE] {
219    fn from(input_voice_audio: InputVoiceAudio) -> Self {
220        let peer_input_audio: PeerInputAudio = input_voice_audio.into();
221        peer_input_audio.into()
222    }
223}
224
225use crate::sync::{broadcast, mpsc};
226use async_trait::async_trait;
227#[async_trait]
228pub trait STTPlugin: Plugin + Send + Sync {
229    async fn start(
230        &self,
231        audio_rx: mpsc::Receiver<InputVoiceAudio>,
232        transcript_tx: mpsc::Sender<SpeechTranscript>,
233        speech_detected: SpeechDetected,
234    ) -> Result<(), String>;
235}
236#[async_trait]
237pub trait TTSPlugin: Plugin + Send + Sync {
238    async fn start(
239        &self,
240        speech_rx: broadcast::Receiver<CognitiveOutputSpeech>,
241        audio_tx: mpsc::Sender<CognitiveOutputAudio>,
242    ) -> Result<(), String>;
243}
244#[async_trait]
245pub trait DiarizationPlugin: Plugin + Send + Sync {
246    async fn start(
247        &self,
248        audio_rx: broadcast::Receiver<InputVoiceAudio>,
249        segment_tx: mpsc::Sender<SpeakerSegment>,
250    ) -> Result<(), String>;
251
252    fn heuristic(&self) -> Option<crate::speech_to_text::SpeakerHeuristicCallback> {
253        None
254    }
255}