1use crate::audio::AudioInput;
25use crate::config::{Config, ValidatedConfig};
26use crate::doctor::{run_doctor, DoctorReport};
27use crate::error::{Result, UserError};
28use crate::observability::{Metrics, MetricsSnapshot};
29use crate::providers::local::{LocalWhisperProvider, SttContextPool};
30use crate::providers::{TranscriptionOptions, TranscriptionProvider, TranscriptionResult};
31use crate::runtime::{GovernorConfig, ResourceGovernor};
32use crate::support::{build_support_bundle, SupportBundle};
33use std::sync::atomic::{AtomicBool, Ordering};
34use std::sync::Arc;
35use std::time::Instant;
36
37#[cfg(feature = "tts")]
38use crate::tts::local::{LocalTtsProvider, TtsSessionPool};
39#[cfg(feature = "tts")]
40use crate::tts::provider::{SynthesisOptions, SynthesisProvider, SynthesisResult};
41
42pub struct AurumEngine {
44 config: ValidatedConfig,
45 governor: Arc<ResourceGovernor>,
46 metrics: Arc<Metrics>,
47 stt_pool: Arc<SttContextPool>,
48 #[cfg(feature = "tts")]
49 tts_pool: Arc<TtsSessionPool>,
50 closed: AtomicBool,
51}
52
53impl std::fmt::Debug for AurumEngine {
54 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
55 let mut d = f.debug_struct("AurumEngine");
56 d.field("config", &self.config)
57 .field("closed", &self.closed.load(Ordering::SeqCst))
58 .field("metrics", &self.metrics.snapshot())
59 .field("stt_resident", &self.stt_pool.resident_len());
60 #[cfg(feature = "tts")]
61 d.field("tts_resident", &self.tts_pool.resident_len());
62 d.finish_non_exhaustive()
63 }
64}
65
66impl AurumEngine {
67 pub fn new(config: ValidatedConfig) -> Self {
69 Self::with_governor(config, GovernorConfig::default())
70 .expect("default GovernorConfig is always valid")
71 }
72
73 pub fn with_governor(config: ValidatedConfig, gov: GovernorConfig) -> Result<Self> {
77 let governor = Arc::new(ResourceGovernor::try_new(gov)?);
78 Ok(Self {
79 config,
80 governor,
81 metrics: Arc::new(Metrics::new()),
82 stt_pool: Arc::new(SttContextPool::new()),
83 #[cfg(feature = "tts")]
84 tts_pool: Arc::new(TtsSessionPool::new()),
85 closed: AtomicBool::new(false),
86 })
87 }
88
89 pub fn load() -> Result<Self> {
91 Ok(Self::new(ValidatedConfig::load()?))
92 }
93
94 pub fn load_from_required(path: &std::path::Path) -> Result<Self> {
96 Ok(Self::new(ValidatedConfig::load_from_required(path)?))
97 }
98
99 pub fn from_config(cfg: Config) -> Result<Self> {
101 Ok(Self::new(ValidatedConfig::try_from_config(cfg)?))
102 }
103
104 pub fn config(&self) -> &Config {
105 self.config.as_ref()
106 }
107
108 pub fn validated_config(&self) -> &ValidatedConfig {
109 &self.config
110 }
111
112 pub fn governor(&self) -> &Arc<ResourceGovernor> {
113 &self.governor
114 }
115
116 pub fn metrics(&self) -> &Arc<Metrics> {
117 &self.metrics
118 }
119
120 pub fn stt_pool(&self) -> &Arc<SttContextPool> {
121 &self.stt_pool
122 }
123
124 #[cfg(feature = "tts")]
125 pub fn tts_pool(&self) -> &Arc<TtsSessionPool> {
126 &self.tts_pool
127 }
128
129 pub fn is_closed(&self) -> bool {
130 self.closed.load(Ordering::SeqCst)
131 }
132
133 fn ensure_open(&self) -> Result<()> {
134 if self.is_closed() {
135 return Err(UserError::Other {
136 message: "AurumEngine is closed".into(),
137 }
138 .into());
139 }
140 Ok(())
141 }
142
143 pub fn local_whisper(&self) -> Result<LocalWhisperProvider> {
145 self.ensure_open()?;
146 Ok(LocalWhisperProvider::with_runtime(
147 self.cache_dir().to_path_buf(),
148 Arc::clone(&self.stt_pool),
149 Arc::clone(&self.governor),
150 )
151 .with_progress(false))
152 }
153
154 #[cfg(feature = "tts")]
156 pub fn local_tts(&self) -> Result<LocalTtsProvider> {
157 self.ensure_open()?;
158 Ok(LocalTtsProvider::with_runtime(
159 self.cache_dir().to_path_buf(),
160 Arc::clone(&self.tts_pool),
161 Arc::clone(&self.governor),
162 )
163 .with_progress(false)
164 .with_max_chars(self.config.as_ref().tts_max_chars))
165 }
166
167 pub async fn transcribe(
169 &self,
170 input: &AudioInput,
171 options: &TranscriptionOptions,
172 ) -> Result<TranscriptionResult> {
173 self.ensure_open()?;
174 self.metrics.record_start();
175 let start = Instant::now();
176 let provider = self.local_whisper()?.with_progress(false);
177 let out = provider.transcribe(input, options).await;
178 match &out {
179 Ok(_) => self.metrics.record_complete(start.elapsed()),
180 Err(_) => self.metrics.record_failed(),
181 }
182 out
183 }
184
185 pub async fn transcribe_pcm(
187 &self,
188 samples: &[f32],
189 options: &TranscriptionOptions,
190 ) -> Result<TranscriptionResult> {
191 self.ensure_open()?;
192 self.metrics.record_start();
193 let start = Instant::now();
194 let provider = self.local_whisper()?.with_progress(false);
195 let out = provider.transcribe_pcm(samples, options).await;
196 match &out {
197 Ok(_) => self.metrics.record_complete(start.elapsed()),
198 Err(_) => self.metrics.record_failed(),
199 }
200 out
201 }
202
203 pub async fn preload_stt(&self, model: &str) -> Result<std::path::PathBuf> {
205 self.ensure_open()?;
206 self.local_whisper()?.preload(model).await
207 }
208
209 #[cfg(feature = "tts")]
211 pub async fn synthesize(
212 &self,
213 text: &str,
214 options: &SynthesisOptions,
215 ) -> Result<SynthesisResult> {
216 self.ensure_open()?;
217 self.metrics.record_start();
218 let start = Instant::now();
219 let provider = self.local_tts()?;
220 let out = provider.synthesize(text, options).await;
221 match &out {
222 Ok(_) => self.metrics.record_complete(start.elapsed()),
223 Err(_) => self.metrics.record_failed(),
224 }
225 out
226 }
227
228 pub fn clear_model_caches(&self) {
230 self.stt_pool.clear();
231 #[cfg(feature = "tts")]
232 self.tts_pool.clear();
233 }
234
235 pub fn shutdown(&self) {
239 self.closed.store(true, Ordering::SeqCst);
240 self.clear_model_caches();
241 }
242
243 pub fn doctor(&self) -> DoctorReport {
245 run_doctor(self.config.as_ref())
246 }
247
248 pub fn support_bundle(&self, user_notes: Option<String>) -> SupportBundle {
250 let mut bundle = build_support_bundle(self.config.as_ref(), user_notes);
251 bundle.metrics = self.metrics.snapshot();
252 bundle.redaction_notes.push(format!(
253 "metrics are engine-local; stt_resident={}{}",
254 self.stt_pool.resident_len(),
255 {
256 #[cfg(feature = "tts")]
257 {
258 format!(", tts_resident={}", self.tts_pool.resident_len())
259 }
260 #[cfg(not(feature = "tts"))]
261 {
262 String::new()
263 }
264 }
265 ));
266 bundle
267 }
268
269 pub fn metrics_snapshot(&self) -> MetricsSnapshot {
270 self.metrics.snapshot()
271 }
272
273 pub fn cache_dir(&self) -> &std::path::Path {
275 &self.config.as_ref().cache_dir
276 }
277}
278
279impl Drop for AurumEngine {
280 fn drop(&mut self) {
281 self.closed.store(true, Ordering::SeqCst);
282 self.clear_model_caches();
283 }
284}
285
286#[cfg(test)]
287mod tests {
288 use super::*;
289
290 #[test]
291 fn independent_engines_have_independent_metrics_and_pools() {
292 let a = AurumEngine::load().unwrap();
293 let b = AurumEngine::load().unwrap();
294 a.metrics().record_start();
295 a.metrics()
296 .record_complete(std::time::Duration::from_millis(1));
297 assert_eq!(a.metrics_snapshot().ops_started, 1);
298 assert_eq!(b.metrics_snapshot().ops_started, 0);
299 assert!(!std::ptr::eq(
300 Arc::as_ptr(a.governor()),
301 Arc::as_ptr(b.governor())
302 ));
303 assert!(!std::ptr::eq(
304 Arc::as_ptr(a.stt_pool()),
305 Arc::as_ptr(b.stt_pool())
306 ));
307 #[cfg(feature = "tts")]
308 assert!(!std::ptr::eq(
309 Arc::as_ptr(a.tts_pool()),
310 Arc::as_ptr(b.tts_pool())
311 ));
312 let process = crate::providers::local::process_global_stt_pool();
314 assert!(!std::ptr::eq(
315 Arc::as_ptr(a.stt_pool()),
316 Arc::as_ptr(&process)
317 ));
318 }
319
320 #[test]
321 fn shutdown_flags_closed_and_rejects_local_whisper() {
322 let e = AurumEngine::load().unwrap();
323 assert!(!e.is_closed());
324 e.shutdown();
325 assert!(e.is_closed());
326 assert!(e.local_whisper().is_err());
327 }
328
329 #[test]
330 fn doctor_and_support_bundle_work() {
331 let e = AurumEngine::load().unwrap();
332 let d = e.doctor();
333 assert!(!d.checks.is_empty());
334 let b = e.support_bundle(None);
335 assert_eq!(b.schema_version, crate::support::SUPPORT_BUNDLE_VERSION);
336 let json = b.to_json_pretty().unwrap();
337 assert!(json.contains("engine-local") || json.contains("stt_resident"));
338 }
339
340 #[test]
341 fn local_whisper_uses_engine_pool() {
342 let e = AurumEngine::load().unwrap();
343 let p = e.local_whisper().unwrap();
344 assert!(std::ptr::eq(
345 Arc::as_ptr(p.pool()),
346 Arc::as_ptr(e.stt_pool())
347 ));
348 assert!(std::ptr::eq(
349 Arc::as_ptr(p.governor()),
350 Arc::as_ptr(e.governor())
351 ));
352 }
353}