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
9pub 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#[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#[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#[derive(Deserialize, Serialize, Default, Clone, Debug, JsonSchema)]
60pub struct Word {
61 pub start_index: Option<u64>,
63 pub end_index: Option<u64>,
65 pub word: String,
67 pub speaker_hint: Option<String>,
69}
70
71#[derive(Deserialize, Serialize, Default, Clone, Debug, JsonSchema)]
73pub struct SpeechTranscript {
74 pub start_index: u64,
76 pub end_index: u64,
78 pub transcript: String,
80 pub words: Vec<Word>,
82}
83
84#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema)]
86pub enum InputVoiceAudio {
87 Voice(PeerInputAudioIndexed),
89 NoVoice(PeerInputAudioIndexed),
91}
92
93#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema)]
95pub struct PeerInputAudioIndexed {
96 pub audio: PeerInputAudio,
98 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}