#![cfg(feature = "documents")]
use std::path::Path;
use anyhow::Context as _;
use xberg::core::mime;
use xberg::embeddings::{EMBEDDING_PRESETS, EmbeddingPreset};
use crate::config::{DocumentsConfig, LlmConfig};
use crate::extract::doc::{DocConfig, FileMapDoc, extract_doc};
use crate::hashing::{self, Hash};
use crate::lance::DocumentRow;
use crate::scanner::EmbedMode;
use crate::store::Store;
#[derive(Debug, Clone)]
pub(crate) struct PendingDocBatch {
pub rel_path: String,
pub chunk_count: usize,
pub embedding_dim: u16,
pub rows: Vec<DocumentRow>,
}
pub(crate) fn preset_dim(name: &str) -> anyhow::Result<u16> {
let preset: &EmbeddingPreset = EMBEDDING_PRESETS
.iter()
.find(|p| p.name == name)
.with_context(|| format!("unknown xberg embedding preset: {name}"))?;
u16::try_from(preset.dimensions)
.with_context(|| format!("preset {name} dimensions {} exceeds u16", preset.dimensions))
}
pub(crate) fn doc_config_from(cfg: &DocumentsConfig, llm: &LlmConfig, embed: bool) -> DocConfig {
DocConfig {
max_characters: cfg.max_characters,
overlap: cfg.overlap,
embedding_preset: Some(cfg.embedding_preset.clone()),
embed,
language: cfg.language.clone(),
keywords: cfg.keywords.clone(),
ner: cfg.ner.clone(),
summarization: cfg.summarization.clone(),
llm: llm.clone(),
embed_max_threads: cfg.embed_max_threads,
}
}
const ARCHIVE_BINARY_EXTENSIONS: &[&str] = &[
"zip", "tar", "gz", "tgz", "bz2", "tbz2", "xz", "txz", "zst", "zstd", "7z", "rar", "lz", "lz4", "lzma", "br", "cab",
"ar", "iso", "dmg", "jar", "war", "ear", "apk", "whl", "egg", "deb", "rpm", "nupkg", "pkg", "msi", "crate",
"so", "dylib", "dll", "a", "o", "obj", "bin", "exe", "wasm", "class", "pyc", "pyo", "pyd", "node", "pack", "idx",
];
const DENY_MIME: &[&str] = &[
"application/zip",
"application/x-tar",
"application/gzip",
"application/x-7z-compressed",
"application/java-archive",
"application/vnd.rar",
"application/wasm",
"application/x-executable",
"application/x-sharedlib",
"application/x-mach-binary",
"application/octet-stream",
"audio/",
"video/",
"font/",
];
fn is_denied_binary_or_archive(abs: &Path, mime_type: &str, cfg: &DocumentsConfig) -> bool {
if let Some(ext) = abs.extension().and_then(|e| e.to_str()) {
let ext_lower = ext.to_ascii_lowercase();
if ARCHIVE_BINARY_EXTENSIONS.contains(&ext_lower.as_str())
|| cfg
.extension_denylist
.iter()
.any(|e| e.eq_ignore_ascii_case(&ext_lower))
{
return true;
}
}
DENY_MIME.iter().any(|entry| matches_mime(entry, mime_type))
}
pub(crate) fn should_extract_document(abs: &Path, cfg: &DocumentsConfig) -> Option<String> {
if !cfg.enabled {
return None;
}
let mime_type = mime::detect_mime_type(abs, false).ok()?;
if is_denied_binary_or_archive(abs, &mime_type, cfg) {
return None;
}
if cfg.mime_allowlist.is_empty() {
return Some(mime_type);
}
let allowed = cfg.mime_allowlist.iter().any(|entry| matches_mime(entry, &mime_type));
if allowed { Some(mime_type) } else { None }
}
fn matches_mime(entry: &str, mime_type: &str) -> bool {
if entry == mime_type {
return true;
}
if let Some(prefix) = entry.strip_suffix('/') {
return mime_type.starts_with(prefix) && mime_type.as_bytes().get(prefix.len()) == Some(&b'/');
}
false
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn extract_and_persist_doc(
store: &Store,
rel: &str,
abs: &Path,
hash: &Hash,
mime_type: &str,
cfg: &DocumentsConfig,
llm: &LlmConfig,
scope: &str,
mode: EmbedMode,
) -> Result<Option<PendingDocBatch>, anyhow::Error> {
let embed = matches!(mode, EmbedMode::Inline) && cfg.embed;
let hex_buf = hashing::hex_buf(hash);
let hash_hex = hashing::hex_str(&hex_buf);
if let Some(cached) = store.read_doc_by_hex(hash_hex).ok().flatten()
&& cached_doc_is_reusable(&cached, cfg, embed)
{
return Ok(Some(pending_from_doc(&cached, rel, scope, cfg, embed)));
}
let doc_config = doc_config_from(cfg, llm, embed);
let doc: FileMapDoc =
extract_doc(abs, Some(mime_type), &doc_config).with_context(|| format!("extract document {rel}"))?;
store
.write_doc(hash, &doc)
.with_context(|| format!("write doc blob for {rel}"))?;
Ok(Some(pending_from_doc(&doc, rel, scope, cfg, embed)))
}
fn cached_doc_is_reusable(cached: &FileMapDoc, cfg: &DocumentsConfig, embed: bool) -> bool {
if !embed || cached.chunks.is_empty() {
return true;
}
let want_dim = preset_dim(&cfg.embedding_preset).ok();
cached.embedding_dim > 0
&& Some(cached.embedding_dim) == want_dim
&& cached
.chunks
.iter()
.all(|c| c.embedding.len() == cached.embedding_dim as usize)
}
fn pending_from_doc(doc: &FileMapDoc, rel: &str, scope: &str, cfg: &DocumentsConfig, embed: bool) -> PendingDocBatch {
let no_rows = PendingDocBatch {
rel_path: rel.to_string(),
chunk_count: doc.chunks.len(),
embedding_dim: doc.embedding_dim,
rows: Vec::new(),
};
if doc.chunks.len() > cfg.max_chunks_per_document {
tracing::warn!(
rel,
chunks = doc.chunks.len(),
cap = cfg.max_chunks_per_document,
"document exceeds max_chunks_per_document; caching blob but skipping vector rows"
);
return no_rows;
}
if !embed || doc.embedding_dim == 0 || doc.chunks.is_empty() {
return no_rows;
}
let rows = build_doc_rows(doc, rel, scope);
PendingDocBatch {
rel_path: rel.to_string(),
chunk_count: rows.len(),
embedding_dim: doc.embedding_dim,
rows,
}
}
fn build_doc_rows(doc: &FileMapDoc, rel: &str, scope: &str) -> Vec<DocumentRow> {
let scope_owned = scope.to_string();
let rel_owned = rel.to_string();
let mime_owned = doc.mime_type.clone();
doc.chunks
.iter()
.enumerate()
.map(|(idx, chunk)| DocumentRow {
scope: scope_owned.clone(),
path: rel_owned.clone(),
chunk_idx: u32::try_from(idx).unwrap_or(u32::MAX),
mime_type: mime_owned.clone(),
text: chunk.text.clone(),
byte_start: chunk.byte_start,
byte_end: chunk.byte_end,
embedding: chunk.embedding.clone(),
})
.collect()
}
pub(crate) fn delete_stale_documents(store: &mut Store, config: &crate::config::Config, scope: &str, stale: &[String]) {
if stale.is_empty() {
return;
}
for path in stale {
store.remove_doc(path);
}
if store.lance.is_none() && !store.lance_dir_exists() {
return;
}
let model = &config.documents.embedding_preset;
let dim = match preset_dim(model) {
Ok(dim) => dim,
Err(error) => {
tracing::warn!(?error, preset = %model, "doc stale purge: unknown preset; skipping lance delete");
return;
}
};
let lance = match store.lance_or_open(dim, model) {
Ok(lance) => lance.clone(),
Err(error) => {
tracing::warn!(?error, "doc stale purge: open LanceStore failed; skipping");
return;
}
};
for path in stale {
let doc_scope = doc_scope_for(path, scope, config);
if let Err(error) = lance.replace_document(doc_scope.as_ref(), path, Vec::new()) {
tracing::warn!(
rel = %path,
?error,
"doc stale purge failed; search_documents may return a removed path"
);
}
}
}
pub(crate) fn flush_document_batches(
store: &mut Store,
scope: &str,
batches: Vec<PendingDocBatch>,
embedding_model: &str,
) -> usize {
let mut inserted = 0usize;
let Some(dim) = batches.iter().find(|b| b.embedding_dim > 0).map(|b| b.embedding_dim) else {
return 0;
};
match preset_dim(embedding_model) {
Ok(expected) if expected != dim => {
tracing::error!(
preset = %embedding_model,
expected,
actual = dim,
"preset/runtime dim mismatch — refusing to write document batch"
);
return 0;
}
Ok(_) => {}
Err(error) => {
tracing::error!(
?error,
preset = %embedding_model,
"unknown embedding preset — refusing to write document batch"
);
return 0;
}
}
let lance = match store.lance_or_open(dim, embedding_model) {
Ok(s) => s.clone(),
Err(error) => {
tracing::error!(?error, "open LanceStore for document batch failed");
return 0;
}
};
for batch in batches {
if batch.rows.is_empty() {
continue;
}
match lance.replace_document(scope, &batch.rel_path, batch.rows) {
Ok(()) => inserted += 1,
Err(error) => {
tracing::warn!(
rel = %batch.rel_path,
?error,
"lance replace_document failed; document search may be incomplete"
);
}
}
}
inserted
}
pub(crate) fn doc_scope_for<'a>(
rel: &str,
default_scope: &'a str,
config: &crate::config::Config,
) -> std::borrow::Cow<'a, str> {
if !rel.starts_with('/') {
return std::borrow::Cow::Borrowed(default_scope);
}
for raw_root in &config.scan.extra_roots {
if let Ok(canonical) = raw_root.canonicalize()
&& let Some(prefix) = canonical.to_str()
&& rel.starts_with(prefix)
{
return std::borrow::Cow::Owned(format!("path:{prefix}"));
}
}
std::borrow::Cow::Owned(format!("path:{rel}"))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn preset_dim_for_balanced_returns_768() {
let dim = preset_dim("balanced").expect("balanced preset");
assert_eq!(dim, 768);
}
#[test]
fn preset_dim_for_unknown_errors() {
let err = preset_dim("does-not-exist").expect_err("unknown preset");
let msg = err.to_string();
assert!(
msg.contains("does-not-exist"),
"error should name the preset; got: {msg}"
);
}
#[test]
fn matches_mime_exact_and_prefix() {
assert!(matches_mime("application/pdf", "application/pdf"));
assert!(matches_mime("image/", "image/png"));
assert!(matches_mime("image/", "image/jpeg"));
assert!(!matches_mime("image/", "video/mp4"));
assert!(!matches_mime("application/pdf", "application/json"));
assert!(!matches_mime("image/", "imageprocessing/x"));
}
#[test]
fn doc_config_from_propagates_language_settings() {
use crate::config::DocLanguageConfig;
let cfg = DocumentsConfig {
language: DocLanguageConfig {
auto_detect: true,
min_confidence: 0.5,
detect_multiple: true,
..Default::default()
},
..Default::default()
};
let doc_cfg = doc_config_from(&cfg, &LlmConfig::default(), cfg.embed);
assert!(doc_cfg.language.auto_detect);
assert_eq!(doc_cfg.language.min_confidence, 0.5);
assert!(doc_cfg.language.detect_multiple);
}
#[test]
fn doc_config_from_propagates_summarization_and_llm() {
use crate::config::{SummarizationConfig, SummarizationStrategy};
let cfg = DocumentsConfig {
summarization: SummarizationConfig {
enabled: true,
strategy: SummarizationStrategy::Abstractive,
max_tokens: Some(150),
},
..Default::default()
};
let llm = LlmConfig {
model: "openai/gpt-4o".to_string(),
..Default::default()
};
let doc_cfg = doc_config_from(&cfg, &llm, cfg.embed);
assert!(doc_cfg.summarization.enabled);
assert_eq!(doc_cfg.summarization.max_tokens, Some(150));
assert_eq!(doc_cfg.llm.model, "openai/gpt-4o");
}
#[test]
fn should_extract_document_respects_disabled_flag() {
let cfg = DocumentsConfig {
enabled: false,
..Default::default()
};
let out = should_extract_document(Path::new("dummy.pdf"), &cfg);
assert!(out.is_none());
}
#[test]
fn should_extract_document_rejects_archives_and_binaries() {
let cfg = DocumentsConfig::default();
for path in [
"vendor/lib.zip",
"target/app.jar",
"dist/bundle.tar.gz",
"build/libfoo.so",
"pkg/module.wasm",
"out/Main.class",
"wheels/pkg-1.0.whl",
"bin/tool.exe",
"obj/thing.o",
] {
assert!(
should_extract_document(Path::new(path), &cfg).is_none(),
"archive/binary must be denied: {path}"
);
}
}
#[test]
fn should_extract_document_allows_real_documents() {
let cfg = DocumentsConfig::default();
for path in ["docs/manual.pdf", "notes/readme.txt", "report.csv"] {
assert!(
should_extract_document(Path::new(path), &cfg).is_some(),
"extractable document must pass: {path}"
);
}
}
#[test]
fn should_extract_document_honors_extension_denylist_override() {
let cfg = DocumentsConfig {
extension_denylist: vec!["pdf".to_string()],
..Default::default()
};
assert!(should_extract_document(Path::new("docs/manual.pdf"), &cfg).is_none());
assert!(should_extract_document(Path::new("vendor/lib.zip"), &cfg).is_none());
}
#[test]
fn images_pass_but_audio_video_denied() {
let cfg = DocumentsConfig::default();
assert!(should_extract_document(Path::new("assets/photo.png"), &cfg).is_some());
assert!(should_extract_document(Path::new("clips/audio.mp3"), &cfg).is_none());
assert!(should_extract_document(Path::new("clips/movie.mp4"), &cfg).is_none());
}
#[test]
fn doc_scope_keeps_default_for_repo_relative_paths() {
let cfg = crate::config::ConfigV1::with_defaults();
let scope = doc_scope_for("docs/manual.pdf", "repo:origin", &cfg);
assert_eq!(scope, "repo:origin");
assert!(matches!(scope, std::borrow::Cow::Borrowed(_)));
}
#[test]
fn doc_scope_namespaces_external_files_under_their_extra_root() {
let ext = tempfile::tempdir().expect("tempdir");
let ext_canonical = std::fs::canonicalize(ext.path()).unwrap();
let mut cfg = crate::config::ConfigV1::with_defaults();
cfg.scan.extra_roots = vec![ext.path().to_path_buf()];
let file_key = ext_canonical.join("pkg/notes.pdf");
let scope = doc_scope_for(file_key.to_str().unwrap(), "repo:origin", &cfg);
assert_eq!(scope, format!("path:{}", ext_canonical.to_str().unwrap()));
}
}