314 lines
12 KiB
Rust
314 lines
12 KiB
Rust
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, Some(limits), |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<()> {
|
|
let change_revision = self.lexical.media_change_revision(
|
|
&path.to_string_lossy(),
|
|
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,
|
|
"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, change_revision)?;
|
|
}
|
|
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,
|
|
change_revision,
|
|
&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,
|
|
change_revision: i64,
|
|
) -> Result<()> {
|
|
let tags = metadata
|
|
.tags
|
|
.iter()
|
|
.map(|(k, v)| format!("{k}: {v}"))
|
|
.collect::<Vec<_>>()
|
|
.join("\n");
|
|
self.lexical.set_media_metadata_at_revision(
|
|
&path.to_string_lossy(),
|
|
&tags,
|
|
metadata.duration_ms,
|
|
identity.bytes,
|
|
identity.modified_ns,
|
|
change_revision,
|
|
)
|
|
}
|
|
|
|
#[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,
|
|
change_revision: i64,
|
|
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, change_revision)?;
|
|
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(())
|
|
}
|
|
}
|