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

use super::pipeline::BuildWork;
use super::*;
use searchi_media::*;
use store::SourceTask;

#[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, work| {
            service.build(generation, work)?;
            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 = work.report.chunks;
                    let identity = MediaIdentity::read(
                        path,
                        if kind == MediaKind::Image {
                            64 * 1024 * 1024
                        } else {
                            limits.max_file_bytes
                        },
                    );
                    let task = SourceTask {
                        key: super::pipeline::task_key(&["media", &path.to_string_lossy()]),
                        path: path.to_string_lossy().into_owned(),
                        source_type: "media".into(),
                        source_hash: identity
                            .as_ref()
                            .map(|i| i.sha256.clone())
                            .unwrap_or_default(),
                        title_hash: String::new(),
                        policy: service.generation_fingerprint.clone(),
                        passage_id: -id,
                    };
                    work.track(&task)?;
                    service.store.stage_task(generation, &task)?;
                    let result = identity.and_then(|identity| {
                        service.build_media_source(
                            generation, id, path, kind, tools, limits, &identity, &task, work,
                        )?;
                        work.media.push(identity);
                        Ok(())
                    });
                    match result {
                        Ok(()) => work.report.media_files += 1,
                        Err(error) => {
                            if let Err(cleanup) =
                                service
                                    .store
                                    .fail_source(generation, &task, &format!("{error:#}"))
                            {
                                return Err(error.context(format!(
                                    "Failed to discard incomplete media source: {cleanup:#}"
                                )));
                            }
                            if error.downcast_ref::<InferenceFailure>().is_some() {
                                return Err(error);
                            }
                            service.store.media_failure(
                                generation,
                                &task.path,
                                &format!("{error:#}"),
                            )?;
                            work.report.chunks = before_chunks;
                            if work.report.failures.len() < 1000 {
                                work.report.failures.push(MediaFailure {
                                    path: task.path,
                                    message: format!("{error:#}").chars().take(2048).collect(),
                                });
                            }
                            work.report.media_failures += 1;
                        }
                    }
                }
            }
            Ok(())
        })
    }

    #[allow(clippy::too_many_arguments)]
    fn build_media_source(
        &mut self,
        generation: i64,
        id: i64,
        path: &Path,
        kind: MediaKind,
        tools: &FfmpegTools,
        limits: &MediaLimits,
        identity: &MediaIdentity,
        task: &SourceTask,
        work: &mut BuildWork,
    ) -> Result<()> {
        if let Some(count) = self
            .store
            .reuse_media_source(generation, task, id, identity)?
        {
            ensure!(
                work.report.chunks + count <= MAX_SCAN_CHUNKS,
                "Media exceeds generation chunk budget"
            );
            if kind != MediaKind::Image {
                // Metadata remains part of music search after path remapping.
                let metadata = probe(tools, path, limits)?;
                self.refresh_media_metadata(path, identity, &metadata)?;
            }
            identity.ensure_unchanged()?;
            work.report.chunks += count;
            work.report.reused_chunks += count;
            work.report.reused_sources += 1;
            return Ok(());
        }
        let before = work.report.chunks;
        self.build_media_file(
            generation,
            id,
            path,
            kind,
            tools,
            limits,
            identity,
            &mut work.report,
        )?;
        self.store
            .complete_source(generation, task, work.report.chunks - before)?;
        work.report.encoded_sources += 1;
        Ok(())
    }

    fn refresh_media_metadata(
        &mut self,
        path: &Path,
        identity: &MediaIdentity,
        metadata: &MediaProbe,
    ) -> Result<()> {
        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,
        )
    }

    #[allow(clippy::too_many_arguments)]
    fn build_media_file(
        &mut self,
        generation: i64,
        id: i64,
        path: &Path,
        kind: MediaKind,
        tools: &FfmpegTools,
        limits: &MediaLimits,
        identity: &MediaIdentity,
        report: &mut EmbeddingReport,
    ) -> Result<()> {
        let mut ordinal = 0;
        if kind == MediaKind::Image {
            validate_image(path)?;
            ensure!(
                report.chunks < MAX_SCAN_CHUNKS,
                "Media exceeds generation chunk budget"
            );
            report.image_inference_calls += 1;
            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)?;
            self.refresh_media_metadata(path, identity, &metadata)?;
            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"
                        );
                        report.image_inference_calls += 1;
                        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)?;
                    report.audio_inference_calls += 1;
                    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(())
    }
}