533 lines
19 KiB
Rust
533 lines
19 KiB
Rust
//! Thin headless adapter around the shared local search engine.
|
||
|
||
use std::path::PathBuf;
|
||
|
||
use anyhow::{bail, Context, Result};
|
||
use clap::{Parser, Subcommand};
|
||
use searchi_core::metrics::{MetricsMode, SearchResponse, TraceCollector};
|
||
use searchi_core::{IndexOptions, SearchIndex, SearchOptions};
|
||
use searchi_service::HybridSearch;
|
||
mod mcp;
|
||
|
||
#[derive(Debug, Parser)]
|
||
#[command(name = "searchi", version, about = "Search your local files offline")]
|
||
struct Cli {
|
||
#[command(subcommand)]
|
||
command: Command,
|
||
}
|
||
|
||
#[derive(Debug, Subcommand)]
|
||
enum Command {
|
||
/// Serve read-only local agent tools over MCP stdio (no HTTP listener).
|
||
Mcp {
|
||
#[arg(long)]
|
||
db: PathBuf,
|
||
#[arg(long, requires = "embedding_worker")]
|
||
model_dir: Option<PathBuf>,
|
||
#[arg(long, requires = "model_dir")]
|
||
embedding_worker: Option<PathBuf>,
|
||
#[arg(long,default_value_t=1,value_parser=clap::value_parser!(u8).range(1..=8))]
|
||
cpu_threads: u8,
|
||
},
|
||
/// Set up the required Google model in the default local cache.
|
||
Setup {
|
||
/// Set up the local speech recognition models instead of Google embeddings.
|
||
#[arg(long)]
|
||
speech: bool,
|
||
/// Override the platform-local model cache directory.
|
||
#[arg(long)]
|
||
model_dir: Option<PathBuf>,
|
||
},
|
||
/// Discover files and refresh an index for one folder.
|
||
Index {
|
||
/// SQLite index file; parent directory must exist.
|
||
#[arg(long)]
|
||
db: PathBuf,
|
||
/// Folder to index recursively.
|
||
#[arg(long)]
|
||
root: PathBuf,
|
||
/// Additional directory names to exclude; repeat this flag as needed.
|
||
#[arg(long)]
|
||
exclude: Vec<String>,
|
||
},
|
||
/// Search filenames and indexed document content.
|
||
Search {
|
||
#[arg(long)]
|
||
db: PathBuf,
|
||
/// Search text; quote phrases containing spaces in your shell.
|
||
#[arg(default_value = "")]
|
||
query: String,
|
||
/// Maximum number of returned results (1–1000).
|
||
#[arg(long, default_value_t = 20, value_parser = clap::value_parser!(u16).range(1..=1000))]
|
||
limit: u16,
|
||
/// Emit a JSON array for scripting.
|
||
#[arg(long)]
|
||
json: bool,
|
||
/// Include basic measured statistics in JSON output.
|
||
#[arg(long)]
|
||
performance: bool,
|
||
/// Include bounded granular stage timings and diagnostic counters.
|
||
#[arg(long)]
|
||
debug_performance: bool,
|
||
/// Restrict results to this folder and descendants.
|
||
#[arg(long)]
|
||
folder: Option<PathBuf>,
|
||
/// Search indexed source-code files only; supports symbols and semantic code queries.
|
||
#[arg(long)]
|
||
code_only: bool,
|
||
/// File extension, with or without the leading dot.
|
||
#[arg(long)]
|
||
extension: Option<String>,
|
||
/// Exact artist tag; unrelated semantic neighbours cannot bypass it.
|
||
#[arg(long)]
|
||
media_artist: Option<String>,
|
||
/// Exact track/media title tag.
|
||
#[arg(long)]
|
||
media_title: Option<String>,
|
||
#[arg(long)]
|
||
min_size: Option<u64>,
|
||
#[arg(long)]
|
||
max_size: Option<u64>,
|
||
/// Inclusive media duration limit in milliseconds.
|
||
#[arg(long)]
|
||
min_duration_ms: Option<u32>,
|
||
#[arg(long)]
|
||
max_duration_ms: Option<u32>,
|
||
/// Inclusive modification time in Unix milliseconds.
|
||
#[arg(long)]
|
||
modified_after_ms: Option<i64>,
|
||
#[arg(long)]
|
||
modified_before_ms: Option<i64>,
|
||
/// Enable hybrid search with explicitly provisioned local model assets.
|
||
#[arg(long, requires = "embedding_worker")]
|
||
model_dir: Option<PathBuf>,
|
||
/// Local native Google embedding worker, required with --model-dir.
|
||
#[arg(long, requires = "model_dir")]
|
||
embedding_worker: Option<PathBuf>,
|
||
#[arg(long, default_value_t = 1, value_parser = clap::value_parser!(u8).range(1..=8))]
|
||
cpu_threads: u8,
|
||
},
|
||
/// Build a persistent semantic generation from already indexed passages.
|
||
Embed {
|
||
#[arg(long)]
|
||
db: PathBuf,
|
||
#[arg(long)]
|
||
model_dir: PathBuf,
|
||
#[arg(long)]
|
||
embedding_worker: PathBuf,
|
||
#[arg(long, default_value_t = 1, value_parser = clap::value_parser!(u8).range(1..=8))]
|
||
cpu_threads: u8,
|
||
/// Standalone decoder; normally discovered beside the installed CLI.
|
||
#[arg(long, requires = "ffprobe")]
|
||
ffmpeg: Option<PathBuf>,
|
||
#[arg(long, requires = "ffmpeg")]
|
||
ffprobe: Option<PathBuf>,
|
||
/// Explicitly build only document vectors, without media.
|
||
#[arg(long)]
|
||
text_only: bool,
|
||
/// Optional bundled local speech worker; recognizes passages without replacing Google.
|
||
#[arg(long, requires = "speech_model_dir", conflicts_with = "text_only")]
|
||
speech_worker: Option<PathBuf>,
|
||
/// Directory containing verified base and Silero speech models.
|
||
#[arg(long, requires = "speech_worker", conflicts_with = "text_only")]
|
||
speech_model_dir: Option<PathBuf>,
|
||
},
|
||
/// Show the indexed file and content counts.
|
||
Stats {
|
||
#[arg(long)]
|
||
db: PathBuf,
|
||
#[arg(long)]
|
||
json: bool,
|
||
},
|
||
}
|
||
|
||
fn main() -> Result<()> {
|
||
match Cli::parse().command {
|
||
Command::Mcp {
|
||
db,
|
||
model_dir,
|
||
embedding_worker,
|
||
cpu_threads,
|
||
} => {
|
||
require_existing_index(&db)?;
|
||
mcp::serve(db, model_dir, embedding_worker, usize::from(cpu_threads))?;
|
||
}
|
||
Command::Setup { model_dir, speech } => {
|
||
if speech {
|
||
let directory = model_dir
|
||
.map(Ok)
|
||
.unwrap_or_else(searchi_assets::speech::default_model_directory)?;
|
||
let progress = std::sync::atomic::AtomicU64::new(0);
|
||
eprintln!("Setting up local speech models; verified cached assets are reused (149 MB maximum download).");
|
||
searchi_assets::speech::provision(&directory, &progress)?;
|
||
println!("Speech models ready: {}", directory.display());
|
||
return Ok(());
|
||
}
|
||
let directory = model_dir
|
||
.map(Ok)
|
||
.unwrap_or_else(searchi_assets::default_model_directory)?;
|
||
let progress = std::sync::atomic::AtomicU64::new(0);
|
||
eprintln!("Setting up Google EmbeddingGemma 2; cached assets are reused (865 MB maximum download).");
|
||
searchi_assets::provision(&directory, &progress)?;
|
||
println!("Google model ready: {}", directory.display());
|
||
}
|
||
Command::Index { db, root, exclude } => {
|
||
let mut options = IndexOptions::default();
|
||
options.excluded_directory_names.extend(exclude);
|
||
let mut index = SearchIndex::open(&db)
|
||
.with_context(|| format!("Cannot open index {}", db.display()))?;
|
||
let report = index
|
||
.index_folder_with_options(&root, &options)
|
||
.with_context(|| format!("Cannot index folder {}", root.display()))?;
|
||
index
|
||
.prewarm()
|
||
.context("Cannot prepare search statistics")?;
|
||
println!(
|
||
"Discovered: {} Indexed: {} Unchanged: {} Removed: {} Skipped: {} Extraction failures: {}",
|
||
report.discovered, report.indexed, report.unchanged, report.removed, report.skipped, report.failures
|
||
);
|
||
}
|
||
Command::Search {
|
||
db,
|
||
query,
|
||
limit,
|
||
json,
|
||
performance,
|
||
debug_performance,
|
||
folder,
|
||
code_only,
|
||
extension,
|
||
media_artist,
|
||
media_title,
|
||
min_size,
|
||
max_size,
|
||
min_duration_ms,
|
||
max_duration_ms,
|
||
modified_after_ms,
|
||
modified_before_ms,
|
||
model_dir,
|
||
embedding_worker,
|
||
cpu_threads,
|
||
} => {
|
||
if query.trim().is_empty() && media_artist.is_none() && media_title.is_none() {
|
||
bail!("Supply search text or an exact media artist/title constraint");
|
||
}
|
||
require_existing_index(&db)?;
|
||
let options = SearchOptions {
|
||
folder,
|
||
code_only,
|
||
extension,
|
||
media_artist,
|
||
media_title,
|
||
min_size,
|
||
max_size,
|
||
min_duration_ms,
|
||
max_duration_ms,
|
||
modified_after_unix_ms: modified_after_ms,
|
||
modified_before_unix_ms: modified_before_ms,
|
||
};
|
||
let mut trace = TraceCollector::new(if debug_performance {
|
||
MetricsMode::Debug
|
||
} else {
|
||
MetricsMode::Basic
|
||
});
|
||
let semantic_assets = model_dir
|
||
.zip(embedding_worker)
|
||
.filter(|_| !query.trim().is_empty());
|
||
let results = if let Some((model, runtime)) = semantic_assets {
|
||
let mut service = trace.timed("cli.session_initialize", None, |_, _| {
|
||
HybridSearch::open_query(&db, model, runtime, usize::from(cpu_threads))
|
||
})?;
|
||
service
|
||
.search_traced(&query, usize::from(limit), &options, &mut trace, None)
|
||
.context("Hybrid search failed")?
|
||
} else {
|
||
let index = trace
|
||
.timed("cli.index_open", None, |_, _| SearchIndex::open(&db))
|
||
.with_context(|| format!("Cannot open index {}", db.display()))?;
|
||
index
|
||
.search_with_options_traced(
|
||
&query,
|
||
usize::from(limit),
|
||
&options,
|
||
&mut trace,
|
||
None,
|
||
)
|
||
.context("Search failed")?
|
||
};
|
||
let response = SearchResponse {
|
||
metrics: trace.finish(results.len()),
|
||
results,
|
||
};
|
||
if json {
|
||
if performance || debug_performance {
|
||
println!("{}", serde_json::to_string_pretty(&response)?);
|
||
} else {
|
||
println!("{}", serde_json::to_string_pretty(&response.results)?);
|
||
}
|
||
} else if response.results.is_empty() {
|
||
println!("No matching files.");
|
||
} else {
|
||
for result in &response.results {
|
||
println!("{} [{}]", result.path, result.match_kind);
|
||
if let Some(source) = &result.source {
|
||
println!(" {} {}", source.kind, source.number);
|
||
}
|
||
if let Some(error) = &result.extraction_error {
|
||
println!(" Text extraction unavailable: {error}");
|
||
}
|
||
if !result.snippet.is_empty() {
|
||
println!(" {}", result.snippet.replace(['\r', '\n'], " "));
|
||
}
|
||
}
|
||
}
|
||
if !json {
|
||
eprintln!(
|
||
"{} results · {:.2} ms (includes CLI initialization)",
|
||
response.metrics.result_count, response.metrics.total_ms
|
||
);
|
||
if debug_performance {
|
||
eprintln!("{}", serde_json::to_string_pretty(&response.metrics)?);
|
||
}
|
||
}
|
||
}
|
||
Command::Embed {
|
||
db,
|
||
model_dir,
|
||
embedding_worker,
|
||
cpu_threads,
|
||
ffmpeg,
|
||
ffprobe,
|
||
text_only,
|
||
speech_worker,
|
||
speech_model_dir,
|
||
} => {
|
||
require_existing_index(&db)?;
|
||
let mut service =
|
||
HybridSearch::open(&db, model_dir, embedding_worker, usize::from(cpu_threads))?;
|
||
if let Some((worker, directory)) = speech_worker.zip(speech_model_dir) {
|
||
let speech = searchi_speech::SpeechWorker::load(
|
||
&worker,
|
||
&directory.join("ggml-base.bin"),
|
||
&directory.join("ggml-silero-v6.2.0.bin"),
|
||
searchi_speech::Device::Auto,
|
||
)?;
|
||
service = service.with_speech_provider(Box::new(speech))?;
|
||
}
|
||
let report = if text_only {
|
||
service.index_embeddings()
|
||
} else {
|
||
let tools = resolve_media_tools(ffmpeg, ffprobe)?;
|
||
service
|
||
.index_embeddings_with_media(&tools, &searchi_service::MediaLimits::default())
|
||
}
|
||
.context("Semantic indexing failed; prior generation remains active")?;
|
||
println!("{}", serde_json::to_string_pretty(&report)?);
|
||
}
|
||
Command::Stats { db, json } => {
|
||
require_existing_index(&db)?;
|
||
let index = SearchIndex::open(&db)
|
||
.with_context(|| format!("Cannot open index {}", db.display()))?;
|
||
let stats = index.stats().context("Cannot read index statistics")?;
|
||
if json {
|
||
println!("{}", serde_json::to_string_pretty(&stats)?);
|
||
} else {
|
||
println!(
|
||
"Files: {} Files with searchable content: {}",
|
||
stats.files, stats.content_files
|
||
);
|
||
}
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
fn resolve_media_tools(
|
||
ffmpeg: Option<PathBuf>,
|
||
ffprobe: Option<PathBuf>,
|
||
) -> Result<searchi_service::FfmpegTools> {
|
||
if let (Some(ffmpeg), Some(ffprobe)) = (ffmpeg, ffprobe) {
|
||
return searchi_service::FfmpegTools::new(ffmpeg, ffprobe);
|
||
}
|
||
if let Some(tools) = searchi_service::FfmpegTools::from_environment()? {
|
||
return Ok(tools);
|
||
}
|
||
let development =
|
||
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../.searchi/media-runtime/bin");
|
||
let directory = if cfg!(debug_assertions) && development.is_dir() {
|
||
development
|
||
} else {
|
||
std::env::current_exe()?
|
||
.parent()
|
||
.context("CLI installation directory is unavailable")?
|
||
.join("media-runtime/bin")
|
||
};
|
||
searchi_service::FfmpegTools::new(
|
||
directory.join(if cfg!(windows) { "ffmpeg.exe" } else { "ffmpeg" }),
|
||
directory.join(if cfg!(windows) { "ffprobe.exe" } else { "ffprobe" }),
|
||
).context("Media tools are missing from this installation; repair the headless package or explicitly use --text-only")
|
||
}
|
||
|
||
/// Read commands must not silently create a new empty database after a typo.
|
||
fn require_existing_index(db: &std::path::Path) -> Result<()> {
|
||
if !db.is_file() {
|
||
bail!(
|
||
"Index {} does not exist; run `searchi index --db <path> --root <folder>` first",
|
||
db.display()
|
||
);
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn speech_indexing_requires_models_and_cannot_downgrade_to_text_only() {
|
||
let base = [
|
||
"searchi",
|
||
"embed",
|
||
"--db",
|
||
"index.db",
|
||
"--model-dir",
|
||
"google",
|
||
"--embedding-worker",
|
||
"google-worker",
|
||
];
|
||
assert!(
|
||
Cli::try_parse_from(base.into_iter().chain(["--speech-worker", "speech-worker"]))
|
||
.is_err()
|
||
);
|
||
assert!(Cli::try_parse_from(base.into_iter().chain([
|
||
"--speech-worker",
|
||
"speech-worker",
|
||
"--speech-model-dir",
|
||
"speech",
|
||
"--text-only"
|
||
]))
|
||
.is_err());
|
||
assert!(Cli::try_parse_from(base.into_iter().chain([
|
||
"--speech-worker",
|
||
"speech-worker",
|
||
"--speech-model-dir",
|
||
"speech"
|
||
]))
|
||
.is_ok());
|
||
assert!(Cli::try_parse_from(["searchi", "setup", "--speech"]).is_ok());
|
||
}
|
||
|
||
#[test]
|
||
fn rejects_zero_or_excessive_result_limits() {
|
||
for limit in ["0", "1001"] {
|
||
assert!(Cli::try_parse_from([
|
||
"searchi", "search", "--db", "index.db", "hello", "--limit", limit
|
||
])
|
||
.is_err());
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn requires_complete_semantic_configuration_and_bounded_cpu_budget() {
|
||
assert!(Cli::try_parse_from([
|
||
"searchi",
|
||
"search",
|
||
"--db",
|
||
"index.db",
|
||
"query",
|
||
"--model-dir",
|
||
"models"
|
||
])
|
||
.is_err());
|
||
assert!(Cli::try_parse_from([
|
||
"searchi",
|
||
"search",
|
||
"--db",
|
||
"index.db",
|
||
"query",
|
||
"--cpu-threads",
|
||
"0"
|
||
])
|
||
.is_err());
|
||
}
|
||
|
||
#[test]
|
||
fn supports_json_search_and_quoted_query() {
|
||
let cli = Cli::try_parse_from([
|
||
"searchi",
|
||
"search",
|
||
"--db",
|
||
"index.db",
|
||
"termination clause",
|
||
"--json",
|
||
])
|
||
.unwrap();
|
||
assert!(matches!(
|
||
cli.command,
|
||
Command::Search {
|
||
json: true,
|
||
limit: 20,
|
||
..
|
||
}
|
||
));
|
||
}
|
||
#[test]
|
||
fn exact_media_cli_filters_match_real_metadata_without_query_or_model() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let file = dir.path().join("track.ogg");
|
||
std::fs::write(&file, b"metadata storage fixture").unwrap();
|
||
let mut index = SearchIndex::open(":memory:").unwrap();
|
||
index.index_folder(dir.path()).unwrap();
|
||
let path = file.canonicalize().unwrap().to_string_lossy().into_owned();
|
||
let (bytes, modified) = index.file_signature(&path).unwrap().unwrap();
|
||
let revision = index.media_change_revision(&path, bytes, modified).unwrap();
|
||
let tags = std::collections::BTreeMap::from([
|
||
("artist".to_owned(), "Björk".to_owned()),
|
||
("title".to_owned(), "Jóga".to_owned()),
|
||
("comment".to_owned(), "artist: Other Artist".to_owned()),
|
||
]);
|
||
index
|
||
.set_structured_media_metadata_at_revision(
|
||
&path, &tags, 1000, bytes, modified, revision,
|
||
)
|
||
.unwrap();
|
||
let cli = Cli::try_parse_from([
|
||
"searchi",
|
||
"search",
|
||
"--db",
|
||
"index.db",
|
||
"--media-artist",
|
||
" BJÖRK ",
|
||
"--media-title",
|
||
"JÓGA",
|
||
])
|
||
.unwrap();
|
||
let Command::Search {
|
||
query,
|
||
media_artist,
|
||
media_title,
|
||
..
|
||
} = cli.command
|
||
else {
|
||
panic!("search command");
|
||
};
|
||
assert!(query.is_empty());
|
||
let mut options = SearchOptions {
|
||
media_artist,
|
||
media_title,
|
||
..Default::default()
|
||
};
|
||
assert_eq!(
|
||
index.search_with_options(&query, 1, &options).unwrap()[0].filename,
|
||
"track.ogg"
|
||
);
|
||
options.media_artist = Some("Other Artist".into());
|
||
assert!(index
|
||
.search_with_options(&query, 1, &options)
|
||
.unwrap()
|
||
.is_empty());
|
||
options.media_artist = Some(" ".into());
|
||
assert!(index.search_with_options(&query, 1, &options).is_err());
|
||
}
|
||
}
|