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 }
71
72 pub fn with_governor(config: ValidatedConfig, gov: GovernorConfig) -> Self {
74 Self {
75 config,
76 governor: Arc::new(ResourceGovernor::new(gov)),
77 metrics: Arc::new(Metrics::new()),
78 stt_pool: Arc::new(SttContextPool::new()),
79 #[cfg(feature = "tts")]
80 tts_pool: Arc::new(TtsSessionPool::new()),
81 closed: AtomicBool::new(false),
82 }
83 }
84
85 pub fn load() -> Result<Self> {
87 Ok(Self::new(ValidatedConfig::load()?))
88 }
89
90 pub fn load_from_required(path: &std::path::Path) -> Result<Self> {
92 Ok(Self::new(ValidatedConfig::load_from_required(path)?))
93 }
94
95 pub fn from_config(cfg: Config) -> Result<Self> {
97 Ok(Self::new(ValidatedConfig::try_from_config(cfg)?))
98 }
99
100 pub fn config(&self) -> &Config {
101 self.config.as_ref()
102 }
103
104 pub fn validated_config(&self) -> &ValidatedConfig {
105 &self.config
106 }
107
108 pub fn governor(&self) -> &Arc<ResourceGovernor> {
109 &self.governor
110 }
111
112 pub fn metrics(&self) -> &Arc<Metrics> {
113 &self.metrics
114 }
115
116 pub fn stt_pool(&self) -> &Arc<SttContextPool> {
117 &self.stt_pool
118 }
119
120 #[cfg(feature = "tts")]
121 pub fn tts_pool(&self) -> &Arc<TtsSessionPool> {
122 &self.tts_pool
123 }
124
125 pub fn is_closed(&self) -> bool {
126 self.closed.load(Ordering::SeqCst)
127 }
128
129 fn ensure_open(&self) -> Result<()> {
130 if self.is_closed() {
131 return Err(UserError::Other {
132 message: "AurumEngine is closed".into(),
133 }
134 .into());
135 }
136 Ok(())
137 }
138
139 pub fn local_whisper(&self) -> Result<LocalWhisperProvider> {
141 self.ensure_open()?;
142 Ok(LocalWhisperProvider::with_runtime(
143 self.cache_dir().to_path_buf(),
144 Arc::clone(&self.stt_pool),
145 Arc::clone(&self.governor),
146 )
147 .with_progress(false))
148 }
149
150 #[cfg(feature = "tts")]
152 pub fn local_tts(&self) -> Result<LocalTtsProvider> {
153 self.ensure_open()?;
154 Ok(LocalTtsProvider::with_runtime(
155 self.cache_dir().to_path_buf(),
156 Arc::clone(&self.tts_pool),
157 Arc::clone(&self.governor),
158 )
159 .with_progress(false)
160 .with_max_chars(self.config.as_ref().tts_max_chars))
161 }
162
163 pub async fn transcribe(
165 &self,
166 input: &AudioInput,
167 options: &TranscriptionOptions,
168 ) -> Result<TranscriptionResult> {
169 self.ensure_open()?;
170 self.metrics.record_start();
171 let start = Instant::now();
172 let provider = self.local_whisper()?.with_progress(false);
173 let out = provider.transcribe(input, options).await;
174 match &out {
175 Ok(_) => self.metrics.record_complete(start.elapsed()),
176 Err(_) => self.metrics.record_failed(),
177 }
178 out
179 }
180
181 pub async fn transcribe_pcm(
183 &self,
184 samples: &[f32],
185 options: &TranscriptionOptions,
186 ) -> Result<TranscriptionResult> {
187 self.ensure_open()?;
188 self.metrics.record_start();
189 let start = Instant::now();
190 let provider = self.local_whisper()?.with_progress(false);
191 let out = provider.transcribe_pcm(samples, options).await;
192 match &out {
193 Ok(_) => self.metrics.record_complete(start.elapsed()),
194 Err(_) => self.metrics.record_failed(),
195 }
196 out
197 }
198
199 pub async fn preload_stt(&self, model: &str) -> Result<std::path::PathBuf> {
201 self.ensure_open()?;
202 self.local_whisper()?.preload(model).await
203 }
204
205 #[cfg(feature = "tts")]
207 pub async fn synthesize(
208 &self,
209 text: &str,
210 options: &SynthesisOptions,
211 ) -> Result<SynthesisResult> {
212 self.ensure_open()?;
213 self.metrics.record_start();
214 let start = Instant::now();
215 let provider = self.local_tts()?;
216 let out = provider.synthesize(text, options).await;
217 match &out {
218 Ok(_) => self.metrics.record_complete(start.elapsed()),
219 Err(_) => self.metrics.record_failed(),
220 }
221 out
222 }
223
224 pub fn clear_model_caches(&self) {
226 self.stt_pool.clear();
227 #[cfg(feature = "tts")]
228 self.tts_pool.clear();
229 }
230
231 pub fn shutdown(&self) {
235 self.closed.store(true, Ordering::SeqCst);
236 self.clear_model_caches();
237 }
238
239 pub fn doctor(&self) -> DoctorReport {
241 run_doctor(self.config.as_ref())
242 }
243
244 pub fn support_bundle(&self, user_notes: Option<String>) -> SupportBundle {
246 let mut bundle = build_support_bundle(self.config.as_ref(), user_notes);
247 bundle.metrics = self.metrics.snapshot();
248 bundle.redaction_notes.push(format!(
249 "metrics are engine-local; stt_resident={}{}",
250 self.stt_pool.resident_len(),
251 {
252 #[cfg(feature = "tts")]
253 {
254 format!(", tts_resident={}", self.tts_pool.resident_len())
255 }
256 #[cfg(not(feature = "tts"))]
257 {
258 String::new()
259 }
260 }
261 ));
262 bundle
263 }
264
265 pub fn metrics_snapshot(&self) -> MetricsSnapshot {
266 self.metrics.snapshot()
267 }
268
269 pub fn cache_dir(&self) -> &std::path::Path {
271 &self.config.as_ref().cache_dir
272 }
273}
274
275impl Drop for AurumEngine {
276 fn drop(&mut self) {
277 self.closed.store(true, Ordering::SeqCst);
278 self.clear_model_caches();
279 }
280}
281
282#[cfg(test)]
283mod tests {
284 use super::*;
285
286 #[test]
287 fn independent_engines_have_independent_metrics_and_pools() {
288 let a = AurumEngine::load().unwrap();
289 let b = AurumEngine::load().unwrap();
290 a.metrics().record_start();
291 a.metrics()
292 .record_complete(std::time::Duration::from_millis(1));
293 assert_eq!(a.metrics_snapshot().ops_started, 1);
294 assert_eq!(b.metrics_snapshot().ops_started, 0);
295 assert!(!std::ptr::eq(
296 Arc::as_ptr(a.governor()),
297 Arc::as_ptr(b.governor())
298 ));
299 assert!(!std::ptr::eq(
300 Arc::as_ptr(a.stt_pool()),
301 Arc::as_ptr(b.stt_pool())
302 ));
303 #[cfg(feature = "tts")]
304 assert!(!std::ptr::eq(
305 Arc::as_ptr(a.tts_pool()),
306 Arc::as_ptr(b.tts_pool())
307 ));
308 let process = crate::providers::local::process_global_stt_pool();
310 assert!(!std::ptr::eq(
311 Arc::as_ptr(a.stt_pool()),
312 Arc::as_ptr(&process)
313 ));
314 }
315
316 #[test]
317 fn shutdown_flags_closed_and_rejects_local_whisper() {
318 let e = AurumEngine::load().unwrap();
319 assert!(!e.is_closed());
320 e.shutdown();
321 assert!(e.is_closed());
322 assert!(e.local_whisper().is_err());
323 }
324
325 #[test]
326 fn doctor_and_support_bundle_work() {
327 let e = AurumEngine::load().unwrap();
328 let d = e.doctor();
329 assert!(!d.checks.is_empty());
330 let b = e.support_bundle(None);
331 assert_eq!(b.schema_version, crate::support::SUPPORT_BUNDLE_VERSION);
332 let json = b.to_json_pretty().unwrap();
333 assert!(json.contains("engine-local") || json.contains("stt_resident"));
334 }
335
336 #[test]
337 fn local_whisper_uses_engine_pool() {
338 let e = AurumEngine::load().unwrap();
339 let p = e.local_whisper().unwrap();
340 assert!(std::ptr::eq(
341 Arc::as_ptr(p.pool()),
342 Arc::as_ptr(e.stt_pool())
343 ));
344 assert!(std::ptr::eq(
345 Arc::as_ptr(p.governor()),
346 Arc::as_ptr(e.governor())
347 ));
348 }
349}