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

use super::*;
use searchi_media::*;

#[derive(Debug)]
struct InferenceFailure(anyhow::Error);
impl std::fmt::Display for InferenceFailure {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Media inference failed: {}", self.0)
    }
}
impl std::error::Error for InferenceFailure {}

impl HybridSearch {
    /// Publish text and actual image/frame/audio vectors in one atomic generation.
    pub fn index_embeddings_with_media(
        &mut self,
        tools: &FfmpegTools,
        limits: &MediaLimits,
    ) -> Result<EmbeddingReport> {
        let mut fingerprint = generation_fingerprint(self.embedder.fingerprint());
        let defaults = MediaLimits::default();
        if limits.max_frames != defaults.max_frames
            || limits.frame_interval_ms != defaults.frame_interval_ms
            || limits.max_scene_frames != defaults.max_scene_frames
            || limits.audio_window_ms != defaults.audio_window_ms
            || limits.audio_overlap_ms != defaults.audio_overlap_ms
            || limits.max_audio_windows != defaults.max_audio_windows
            || limits.max_duration_ms != defaults.max_duration_ms
            || limits.max_file_bytes != defaults.max_file_bytes
        {
            fingerprint.push_str(&format!(
                ":custom-{}-{}-{}-{}-{}-{}-{}-{}",
                limits.max_frames,
                limits.frame_interval_ms,
                limits.max_scene_frames,
                limits.audio_window_ms,
                limits.audio_overlap_ms,
                limits.max_audio_windows,
                limits.max_duration_ms,
                limits.max_file_bytes
            ));
        }
        self.rebuild_generation(fingerprint, |service, generation| {
            let mut report = service.build(generation)?;
            let mut after = 0;
            let mut media_count = 0usize;
            loop {
                let files = service.store.media_batch(after)?;
                if files.is_empty() {
                    break;
                }
                for (id, path, _, _) in files {
                    after = id;
                    let path = Path::new(&path);
                    let Some(kind) = classify(path) else {
                        continue;
                    };
                    media_count += 1;
                    ensure!(
                        media_count <= 10000,
                        "Media generation exceeds 10000-file resource budget"
                    );
                    let before_chunks = report.chunks;
                    match service.build_media_file(
                        generation,
                        id,
                        path,
                        kind,
                        tools,
                        limits,
                        &mut report,
                    ) {
                        Ok(()) => report.media_files += 1,
                        Err(error) => {
                            if error.downcast_ref::<InferenceFailure>().is_some() {
                                return Err(error);
                            }
                            service.store.media_failure(
                                generation,
                                &path.to_string_lossy(),
                                &format!("{error:#}"),
                            )?;
                            report.chunks = before_chunks;
                            if report.failures.len() < 1000 {
                                report.failures.push(MediaFailure {
                                    path: path.to_string_lossy().into_owned(),
                                    message: format!("{error:#}").chars().take(2048).collect(),
                                });
                            }
                            report.media_failures += 1;
                        }
                    }
                }
            }
            Ok(report)
        })
    }

    #[allow(clippy::too_many_arguments)]
    fn build_media_file(
        &mut self,
        generation: i64,
        id: i64,
        path: &Path,
        kind: MediaKind,
        tools: &FfmpegTools,
        limits: &MediaLimits,
        report: &mut EmbeddingReport,
    ) -> Result<()> {
        let identity = MediaIdentity::read(
            path,
            if kind == MediaKind::Image {
                64 * 1024 * 1024
            } else {
                limits.max_file_bytes
            },
        )?;
        let mut ordinal = 0;
        if kind == MediaKind::Image {
            validate_image(path)?;
            ensure!(
                report.chunks < MAX_SCAN_CHUNKS,
                "Media exceeds generation chunk budget"
            );
            let vectors = self
                .embedder
                .embed_images(&[path], 280)
                .map_err(|e| anyhow::Error::new(InferenceFailure(e)))?;
            ensure!(vectors.len() == 1, "Image embedding count mismatch");
            self.store
                .insert_media(generation, id, ordinal, &identity, "image", 0, &vectors[0])?;
            report.chunks += 1;
        } else {
            let metadata = probe(tools, path, limits)?;
            let tags = metadata
                .tags
                .iter()
                .map(|(k, v)| format!("{k}: {v}"))
                .collect::<Vec<_>>()
                .join("\n");
            self.lexical.set_media_metadata(
                &path.to_string_lossy(),
                &tags,
                metadata.duration_ms,
                identity.bytes,
                identity.modified_ns,
            )?;
            if metadata.has_video && kind == MediaKind::Video {
                let mut seen = std::collections::HashSet::new();
                let mut scenes = 0;
                for time in frame_times(&metadata, limits)? {
                    let frames = decode_frames(
                        tools,
                        path,
                        &metadata,
                        time,
                        scenes < limits.max_scene_frames,
                        limits,
                    )?;
                    for (i, frame) in frames.frames.iter().enumerate() {
                        if !seen.insert(frame.timestamp_ms) {
                            continue;
                        }
                        if i > 0 {
                            scenes += 1;
                        }
                        ensure!(
                            report.chunks < MAX_SCAN_CHUNKS,
                            "Media exceeds generation chunk budget"
                        );
                        let vectors = self
                            .embedder
                            .embed_images(&[frame.path.as_path()], 70)
                            .map_err(|e| anyhow::Error::new(InferenceFailure(e)))?;
                        ensure!(vectors.len() == 1, "Frame embedding count mismatch");
                        self.store.insert_media(
                            generation,
                            id,
                            ordinal,
                            &identity,
                            "video",
                            frame.timestamp_ms,
                            &vectors[0],
                        )?;
                        ordinal += 1;
                        report.chunks += 1;
                    }
                }
            }
            if metadata.has_audio {
                for window in audio_windows(&metadata, limits)? {
                    ensure!(
                        report.chunks < MAX_SCAN_CHUNKS,
                        "Audio exceeds generation chunk budget"
                    );
                    let pcm = decode_audio(tools, path, window, limits)?;
                    let vector = self
                        .embedder
                        .embed_audio(&pcm)
                        .map_err(|e| anyhow::Error::new(InferenceFailure(e)))?;
                    self.store.insert_media(
                        generation,
                        id,
                        ordinal,
                        &identity,
                        "audio",
                        window.start_ms,
                        &vector,
                    )?;
                    ordinal += 1;
                    report.chunks += 1;
                }
            }
        }
        identity.ensure_unchanged()?;
        Ok(())
    }
}