searchi/crates/searchi-service/src/speech.rs

//! Optional local speech evidence. Google remains the only embedding provider.
use super::*;
use searchi_core::SourceLocation;
use searchi_media::{decode_audio_aligned, probe, AudioWindow, MediaIdentity};
use store::SourceTask;

const MAX_SPEECH_WINDOWS: usize = 1024;
const MAX_SEGMENTS: usize = 64;
const MAX_SEGMENT_BYTES: usize = 4096;
const MAX_WINDOW_TEXT: usize = 64 * 1024;

/// Stable automatic-retry identity; excludes generation IDs and metadata epochs.
pub fn speech_policy_key(provider_fingerprint: &str, limits: &MediaLimits) -> Result<String> {
    ensure!(
        !provider_fingerprint.is_empty() && provider_fingerprint.len() <= 4096,
        "Invalid speech provider fingerprint"
    );
    Ok(format!("speech-v1-{}",store::content_hash(&format!("{provider_fingerprint};speech-window=30000;overlap=2000;max-windows=1024;aligned-decoder-v1;max-duration={};max-bytes={}",limits.max_duration_ms,limits.max_file_bytes))))
}

pub trait SpeechProvider: Send {
    /// Includes actual ASR/VAD/alignment assets and policy, not a display name alone.
    fn fingerprint(&self) -> &str;
    fn admission_us(&self) -> u64 {
        0
    }
    /// Starts a new source without retrying a failed window from the previous source.
    fn begin_source(&mut self) -> Result<()> {
        Ok(())
    }
    fn transcribe(&mut self, pcm: &[f32], sample_rate: u32) -> Result<SpeechTranscription>;
}

impl SpeechProvider for searchi_speech::SpeechWorker {
    fn fingerprint(&self) -> &str {
        searchi_speech::SpeechWorker::fingerprint(self)
    }
    fn begin_source(&mut self) -> Result<()> {
        searchi_speech::SpeechWorker::begin_source(self)
    }
    fn transcribe(&mut self, pcm: &[f32], sample_rate: u32) -> Result<SpeechTranscription> {
        searchi_speech::SpeechWorker::transcribe(self, pcm, sample_rate)
    }
}

/// One background admission per bounded window; never claims native preemption.
pub struct BudgetedSpeechProvider {
    inner: Box<dyn SpeechProvider>,
    budget: InferenceBudget,
    admission_us: u64,
}
impl BudgetedSpeechProvider {
    pub fn new(inner: Box<dyn SpeechProvider>, budget: InferenceBudget) -> Self {
        Self {
            inner,
            budget,
            admission_us: 0,
        }
    }
}
impl SpeechProvider for BudgetedSpeechProvider {
    fn fingerprint(&self) -> &str {
        self.inner.fingerprint()
    }
    fn admission_us(&self) -> u64 {
        self.admission_us
    }
    fn begin_source(&mut self) -> Result<()> {
        let start = std::time::Instant::now();
        let permit = self.budget.acquire(InferenceRole::Background, 1);
        self.admission_us = start.elapsed().as_micros().min(u64::MAX as u128) as u64;
        let _permit = permit?;
        self.inner.begin_source()
    }
    fn transcribe(&mut self, pcm: &[f32], sample_rate: u32) -> Result<SpeechTranscription> {
        let start = std::time::Instant::now();
        let permit = self.budget.acquire(InferenceRole::Background, 1);
        self.admission_us = start.elapsed().as_micros().min(u64::MAX as u128) as u64;
        let _permit = permit?;
        self.inner.transcribe(pcm, sample_rate)
    }
}

#[derive(Debug, Serialize)]
pub struct SpeechExcerpt {
    pub path: String,
    pub filename: String,
    pub source: SourceLocation,
    pub end_ms: u32,
    pub text: String,
    pub provenance: &'static str,
}

pub(super) fn speech_task(
    path: &str,
    id: i64,
    identity: Option<&MediaIdentity>,
    policy: &str,
) -> Result<SourceTask> {
    let filename = Path::new(path)
        .file_name()
        .unwrap_or_default()
        .to_string_lossy();
    Ok(SourceTask {
        key: super::pipeline::task_key(&["speech", path]),
        path: path.into(),
        source_type: "speech".into(),
        source_hash: identity
            .map(|identity| identity.sha256.clone())
            .unwrap_or_default(),
        title_hash: store::content_hash(bounded_title(&filename)),
        policy: speech_component_policy(policy),
        passage_id: store::media_passage_id(id, "speech")?,
    })
}

fn speech_windows(start_ms: u32, end_ms: u32) -> Result<Vec<AudioWindow>> {
    ensure!(start_ms <= end_ms, "Invalid speech audio range");
    let mut windows = Vec::new();
    let mut start = start_ms;
    while start < end_ms {
        ensure!(
            windows.len() < MAX_SPEECH_WINDOWS,
            "Speech exceeds 1024 contiguous windows; coverage is incomplete"
        );
        windows.push(AudioWindow {
            start_ms: start,
            duration_ms: (end_ms - start).min(30000),
        });
        if end_ms - start <= 30000 {
            break;
        }
        start = start
            .checked_add(28000)
            .ok_or_else(|| anyhow::anyhow!("Speech timestamp overflow"))?;
    }
    Ok(windows)
}

fn validate_transcription(result: &SpeechTranscription, duration_ms: u32) -> Result<()> {
    ensure!(
        result.passages.len() <= MAX_SEGMENTS,
        "Speech window exceeds 64 passages"
    );
    let mut total = 0usize;
    for passage in &result.passages {
        ensure!(
            passage.start_ms <= passage.end_ms && passage.end_ms <= duration_ms,
            "Speech timestamp lies outside decoded window"
        );
        ensure!(
            !passage.text.trim().is_empty() && passage.text.len() <= MAX_SEGMENT_BYTES,
            "Speech passage text is empty or oversized"
        );
        total = total
            .checked_add(passage.text.len())
            .ok_or_else(|| anyhow::anyhow!("Speech text budget overflow"))?;
        ensure!(
            total <= MAX_WINDOW_TEXT,
            "Speech window exceeds 64 KiB text"
        );
    }
    for value in [
        result.metrics.vad_ms,
        result.metrics.asr_ms,
        result.metrics.load_ms,
        result.metrics.transport_ms,
        result.metrics.fallback_ms,
    ] {
        ensure!(
            value.is_finite() && value >= 0. && value <= (u64::MAX / 1000) as f64,
            "Invalid speech timing metric"
        );
    }
    Ok(())
}

impl HybridSearch {
    pub fn with_speech_provider(mut self, provider: Box<dyn SpeechProvider>) -> Result<Self> {
        ensure!(
            !provider.fingerprint().is_empty() && provider.fingerprint().len() <= 4096,
            "Invalid speech provider fingerprint"
        );
        self.speech = Some(provider);
        Ok(self)
    }

    /// Only returns active, byte-verified indexed speech. Caller may further bind a source handle.
    pub fn speech_excerpt(
        &self,
        path: &str,
        start_ms: u32,
        max_chars: usize,
    ) -> Result<Option<SpeechExcerpt>> {
        self.store.speech_excerpt(path, start_ms, max_chars)
    }

    /// Bind issued search evidence when multiple recognized passages share a timestamp.
    pub fn speech_excerpt_matching(
        &self,
        path: &str,
        start_ms: u32,
        expected_text: &str,
        max_chars: usize,
    ) -> Result<Option<SpeechExcerpt>> {
        self.store
            .speech_excerpt_matching(path, start_ms, Some(expected_text), max_chars)
    }

    #[allow(clippy::too_many_arguments)]
    pub(super) fn build_speech_source(
        &mut self,
        generation: i64,
        id: i64,
        path: &Path,
        tools: &FfmpegTools,
        limits: &MediaLimits,
        identity: &MediaIdentity,
        work: &mut super::pipeline::BuildWork,
    ) -> Result<()> {
        let task = speech_task(
            &path.to_string_lossy(),
            id,
            Some(identity),
            &self.generation_fingerprint,
        )?;
        work.track(&task)?;
        self.store.stage_task(generation, &task)?;
        let before = work.report.chunks;
        let before_speech = work.report.speech_chunks;
        let result =
            self.encode_speech_source(generation, id, path, tools, limits, identity, &task, work);
        match result {
            Ok(()) => Ok(()),
            Err(error) => {
                self.store
                    .fail_source(generation, &task, &format!("{error:#}"))?;
                work.report.chunks = before;
                work.report.speech_chunks = before_speech;
                if error
                    .downcast_ref::<super::media::InferenceFailure>()
                    .is_some()
                {
                    return Err(error);
                }
                if work.report.speech_failures.len() < 1000 {
                    work.report.speech_failures.push(MediaFailure {
                        path: task.path,
                        message: format!("{error:#}").chars().take(2048).collect(),
                    });
                }
                Ok(())
            }
        }
    }

    #[allow(clippy::too_many_arguments)]
    fn encode_speech_source(
        &mut self,
        generation: i64,
        id: i64,
        path: &Path,
        tools: &FfmpegTools,
        limits: &MediaLimits,
        identity: &MediaIdentity,
        task: &SourceTask,
        work: &mut super::pipeline::BuildWork,
    ) -> Result<()> {
        let revision =
            self.lexical
                .media_change_revision(&task.path, identity.bytes, identity.modified_ns)?;
        if let Some(count) = self
            .store
            .reuse_media_source(generation, task, id, identity)?
        {
            ensure!(
                work.report.chunks + count <= MAX_SCAN_CHUNKS,
                "Speech exceeds generation chunk budget"
            );
            identity.ensure_unchanged()?;
            work.report.chunks += count;
            work.report.speech_chunks += count;
            work.report.reused_chunks += count;
            work.report.reused_sources += 1;
            work.report.speech_reused_sources += 1;
            return Ok(());
        }
        let metadata = probe(tools, path, limits)?;
        let filename = path.file_name().unwrap_or_default().to_string_lossy();
        let mut ordinal = 0usize;
        if metadata.has_audio {
            self.speech
                .as_mut()
                .ok_or_else(|| anyhow::anyhow!("Speech provider unavailable"))?
                .begin_source()?;
            for window in speech_windows(metadata.audio_start_ms, metadata.audio_end_ms)? {
                let pcm = decode_audio_aligned(tools, path, window, limits)?;
                let provider = self
                    .speech
                    .as_mut()
                    .ok_or_else(|| anyhow::anyhow!("Speech provider unavailable"))?;
                work.report.speech_windows += 1;
                let result = provider.transcribe(&pcm, 16000);
                let admission = work
                    .report
                    .speech_metrics
                    .entry("admission_us".into())
                    .or_default();
                *admission = admission.saturating_add(provider.admission_us());
                let transcription = result?;
                validate_transcription(&transcription, window.duration_ms)?;
                for (key, value) in [
                    ("vad_us", transcription.metrics.vad_ms),
                    ("asr_us", transcription.metrics.asr_ms),
                    ("load_us", transcription.metrics.load_ms),
                    ("transport_us", transcription.metrics.transport_ms),
                    ("fallback_us", transcription.metrics.fallback_ms),
                ] {
                    let total = work.report.speech_metrics.entry(key.into()).or_default();
                    *total = total.saturating_add((value * 1000.) as u64);
                }
                for passage in transcription.passages {
                    ensure!(
                        work.report.chunks < MAX_SCAN_CHUNKS,
                        "Speech exceeds generation chunk budget"
                    );
                    let start = window
                        .start_ms
                        .checked_add(passage.start_ms)
                        .ok_or_else(|| anyhow::anyhow!("Speech timestamp overflow"))?;
                    let end = window
                        .start_ms
                        .checked_add(passage.end_ms)
                        .ok_or_else(|| anyhow::anyhow!("Speech timestamp overflow"))?;
                    let windows = passage_windows(
                        &mut *self.embedder,
                        &passage.text,
                        bounded_title(&filename),
                    )?;
                    for text_window in windows {
                        ensure!(
                            work.report.chunks < MAX_SCAN_CHUNKS,
                            "Speech exceeds generation chunk budget"
                        );
                        let calls = work
                            .report
                            .speech_metrics
                            .entry("google_embedding_calls".into())
                            .or_default();
                        *calls = calls.saturating_add(1);
                        let vectors = self
                            .embedder
                            .embed_documents(&[&text_window.text], bounded_title(&filename))
                            .map_err(|error| {
                                anyhow::Error::new(super::media::InferenceFailure(error))
                            })?;
                        ensure!(vectors.len() == 1, "Speech embedding count mismatch");
                        self.store.insert_speech(
                            generation,
                            id,
                            ordinal,
                            identity,
                            revision,
                            start,
                            end,
                            &text_window.text,
                            &vectors[0],
                        )?;
                        ordinal += 1;
                        work.report.chunks += 1;
                        work.report.speech_chunks += 1;
                    }
                }
            }
        }
        identity.ensure_unchanged()?;
        ensure!(
            self.lexical
                .media_change_revision(&task.path, identity.bytes, identity.modified_ns)?
                == revision,
            "Speech source changed during recognition"
        );
        self.store.complete_source(generation, task, ordinal)?;
        work.report.encoded_sources += 1;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn new_source_recovers_provider_without_retrying_failed_window() -> Result<()> {
        struct SourceProvider {
            sources: usize,
            calls: usize,
            failed: bool,
        }
        impl SpeechProvider for SourceProvider {
            fn fingerprint(&self) -> &str {
                "lifecycle-test"
            }
            fn begin_source(&mut self) -> Result<()> {
                self.sources += 1;
                self.failed = false;
                Ok(())
            }
            fn transcribe(&mut self, _: &[f32], _: u32) -> Result<SpeechTranscription> {
                self.calls += 1;
                if self.sources == 1 || self.failed {
                    self.failed = true;
                    anyhow::bail!("first source transport failure");
                }
                Ok(SpeechTranscription {
                    passages: Vec::new(),
                    metrics: SpeechMetrics::default(),
                })
            }
        }
        let mut provider = SourceProvider {
            sources: 0,
            calls: 0,
            failed: false,
        };
        provider.begin_source()?;
        let first_failure = provider.transcribe(&[0.], 16000).unwrap_err().to_string();
        assert_eq!(provider.calls, 1, "failed PCM is never retried");
        provider.begin_source()?;
        provider.transcribe(&[0.], 16000)?;
        assert_eq!(provider.calls, 2);
        assert_eq!(provider.sources, 2);
        assert_eq!(first_failure, "first source transport failure");
        Ok(())
    }
    #[test]
    fn point_spans_are_legal_and_invalid_native_evidence_is_rejected() -> Result<()> {
        let mut result = SpeechTranscription {
            passages: vec![SpeechPassage {
                start_ms: 0,
                end_ms: 0,
                text: "Recognized point".into(),
            }],
            metrics: SpeechMetrics::default(),
        };
        validate_transcription(&result, 30000)?;
        result.passages[0].end_ms = 30001;
        assert!(validate_transcription(&result, 30000).is_err());
        result.passages[0].end_ms = 0;
        result.passages[0].text = "x".repeat(4097);
        assert!(validate_transcription(&result, 30000).is_err());
        result.passages[0].text = "valid".into();
        result.metrics.asr_ms = f64::NAN;
        assert!(validate_transcription(&result, 30000).is_err());
        result.metrics.asr_ms = 0.;
        result.metrics.fallback_ms = f64::INFINITY;
        assert!(validate_transcription(&result, 30000).is_err());
        result.metrics.fallback_ms = 12.5;
        validate_transcription(&result, 30000)?;
        Ok(())
    }

    #[test]
    fn contiguous_windows_cover_full_four_hours_and_overflow_is_explicit() -> Result<()> {
        let windows = speech_windows(0, 14400000)?;
        assert_eq!(windows[0].start_ms, 0);
        assert_eq!(
            windows.last().unwrap().start_ms + windows.last().unwrap().duration_ms,
            14400000
        );
        assert!(windows
            .windows(2)
            .all(|pair| pair[1].start_ms == pair[0].start_ms + 28000));
        assert!(speech_windows(0, u32::MAX).is_err());
        let offset = speech_windows(2000, 34000)?;
        assert_eq!(offset[0].start_ms, 2000);
        assert_eq!(
            offset.last().unwrap().start_ms + offset.last().unwrap().duration_ms,
            34000
        );
        Ok(())
    }
}