karl/searchi
1
0 Fork 0
Code Issues Pull Requests Packages Projects Releases 10 Wiki Activity

feat: integrate Google semantic search and branded desktop launcher
Core CI / core (ubuntu-latest) (push) Failing after 3m27s
Core CI / core (macos-latest) (push) Canceled after 0s
Core CI / core (windows-latest) (push) Canceled after 0s
Core CI / desktop (ubuntu-latest) (push) Successful in 5m45s
Core CI / desktop (macos-latest) (push) Canceled after 0s
Core CI / desktop (windows-latest) (push) Canceled after 0s

Browse Source
This commit is contained in:
karl committed 2026-10-12 17:44:56 +02:00
1 parent a90df15503
commit caf6a2f82e
80 files changed
+14971 -837
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/searchi-assets/Cargo.toml
+15
View File
@@ -0,0 +1,15 @@
[package]
name = "searchi-assets"
version.workspace = true
edition.workspace = true
license.workspace = true
description = "Explicit verified local model provisioning shared by desktop and CLI"
[dependencies]
anyhow.workspace = true
serde.workspace = true
serde_json.workspace = true
sha2 = "0.10"
directories = "6"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "rustls-tls"] }
tempfile = "3"
crates/searchi-assets/README.md
+52
View File
@@ -0,0 +1,52 @@
# Searchi assets
GUI-free model provisioning shared by desktop setup and headless CLI setup.
Search never invokes this downloader automatically. Inference and GPU providers
are outside this crate; provisioning tests run on the Windows VM using CPU only.
The compiled [model manifest](../../tools/embeddinggemma2/model-manifest.json)
pins Google EmbeddingGemma 2's text GGUF and multimedia projection to one
Hugging Face export commit, exact byte lengths, and SHA-256 digests. The two
artifacts total 864,676,480 bytes. Source URLs cannot be overridden through the
public API. Executable embedding workers remain bundled release resources and
are never downloaded by this crate.
```rust,no_run
use std::sync::atomic::AtomicU64;
let directory = searchi_assets::default_model_directory()?;
let progress = AtomicU64::new(0);
if !searchi_assets::verify_cached(&directory)? {
// Only after the user explicitly chooses setup/download.
searchi_assets::provision(&directory, &progress)?;
}
# Ok::<(), anyhow::Error>(())
```
- `default_model_directory()` computes the platform-local application data path
via `directories::ProjectDirs`, then appends `models/embeddinggemma2`.
- `verify_cached(directory)` performs offline size and full-digest validation.
It does not trust a same-size file or create directories.
- `provision(directory, progress)` verifies/reuses cached artifacts and downloads
only missing or corrupt files. Call it on a blocking setup worker.
- `TOTAL_BYTES` supplies the progress total to GUI and CLI consumers.
New files are streamed through a 64 KiB buffer and rejected immediately if the
body exceeds the pinned size. Content-Length, when supplied, must also match.
The final byte count and SHA-256 must match before the temporary file is synced
and atomically replaces the destination in the same directory. A failed attempt
leaves any existing destination intact and deletes its temporary file through
RAII. Unix directory metadata is also synced after replacement.
There are three attempts per artifact, with 250/500 ms retry delays. Connections
time out after 20 seconds, and an entire request including streamed body reads
after 900 seconds. Progress includes completed verified artifacts and the current
attempt's received bytes; it resets to completed verified bytes after failure.
Hard process termination can leave orphan `.part` files, which are never treated
as valid models. Downloads resume by reusing completed verified artifacts rather
than trusting partial files.
Tests use a private local HTTP fixture and small authentic byte payloads; they
do not download the production model or rely on a public service. They verify
same-size corruption, cache reuse without URL requests, interrupted-response
retry, streaming size bounds, incorrect digests, atomic replacement, and partial
file cleanup. Linux/macOS runtime checks remain pending.
crates/searchi-assets/src/lib.rs
+520
View File
@@ -0,0 +1,520 @@
//! Explicit model setup shared by the GUI and headless CLI. These functions do
//! not configure an inference provider or download executable workers. Search
//! callers must never call provisioning implicitly.
use anyhow::{bail, Context, Result};
use directories::ProjectDirs;
use reqwest::blocking::Client;
use serde::Deserialize;
use sha2::{Digest, Sha256};
use std::{
fs::{self, File},
io::{BufReader, Read, Write},
path::{Component, Path, PathBuf},
sync::atomic::{AtomicU64, Ordering},
time::Duration,
};
/// Exact bytes of the pinned text model plus its multimedia projection.
pub const TOTAL_BYTES: u64 = 864_676_480;
const ATTEMPTS: usize = 3;
#[derive(Deserialize)]
struct Manifest {
model_id: String,
export_repository: String,
revision: String,
files: Vec<Artifact>,
}
#[derive(Deserialize)]
struct Artifact {
path: String,
bytes: u64,
sha256: String,
}
fn manifest() -> Result<Manifest> {
let manifest: Manifest = serde_json::from_str(include_str!(
"../../../tools/embeddinggemma2/model-manifest.json"
))?;
if manifest.model_id != "google/embeddinggemma-2"
|| manifest.export_repository != "ggml-org/embeddinggemma-2-GGUF"
|| manifest.revision.len() != 40
|| !manifest
.revision
.bytes()
.all(|byte| byte.is_ascii_hexdigit())
|| manifest.files.len() != 2
|| manifest
.files
.iter()
.map(|artifact| artifact.bytes)
.sum::<u64>()
!= TOTAL_BYTES
{
bail!("invalid bundled model manifest");
}
for artifact in &manifest.files {
validate_artifact(artifact)?;
if !matches!(
artifact.path.as_str(),
"embeddinggemma-2-Q8_0.gguf" | "mmproj-embeddinggemma-2-Q8_0.gguf"
) {
bail!("bundled model manifest contains an unexpected artifact");
}
}
if manifest.files[0].path == manifest.files[1].path {
bail!("duplicate model manifest artifact");
}
Ok(manifest)
}
fn validate_artifact(artifact: &Artifact) -> Result<()> {
let mut components = Path::new(&artifact.path).components();
if !matches!(components.next(), Some(Component::Normal(_)))
|| components.next().is_some()
|| artifact.path.contains(['/', '\\'])
|| artifact.bytes == 0
|| artifact.sha256.len() != 64
|| !artifact.sha256.bytes().all(|byte| byte.is_ascii_hexdigit())
{
bail!("invalid artifact path, size, or digest");
}
Ok(())
}
/// Platform-local application data, consistently used by CLI and desktop.
/// This computes the path without creating files or accessing the network.
pub fn default_model_directory() -> Result<PathBuf> {
let project = ProjectDirs::from("dev", "Searchi", "Searchi")
.context("platform application data directory is unavailable")?;
Ok(project
.data_local_dir()
.join("models")
.join("embeddinggemma2"))
}
/// Offline validation includes file length and full SHA-256, not mere existence.
pub fn verify_cached(directory: &Path) -> Result<bool> {
for artifact in manifest()?.files {
if !verified(&directory.join(&artifact.path), &artifact)? {
return Ok(false);
}
}
Ok(true)
}
fn verified(path: &Path, artifact: &Artifact) -> Result<bool> {
let file = match File::open(path) {
Ok(file) => file,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false),
Err(error) => return Err(error).context("open cached model artifact"),
};
if file.metadata()?.len() != artifact.bytes {
return Ok(false);
}
let mut reader = BufReader::new(file).take(artifact.bytes + 1);
let mut hash = Sha256::new();
let mut size = 0u64;
let mut buffer = [0u8; 65536];
loop {
let read = reader.read(&mut buffer)?;
if read == 0 {
break;
}
size += read as u64;
if size > artifact.bytes {
return Ok(false);
}
hash.update(&buffer[..read]);
}
Ok(size == artifact.bytes
&& format!("{:x}", hash.finalize()).eq_ignore_ascii_case(&artifact.sha256))
}
/// Explicit blocking setup: pinned HTTPS artifacts, three attempts per missing
/// or corrupt artifact, streaming size/SHA checks, atomic replacement. Progress
/// is completed verified bytes plus bytes of the current attempt. On failure it
/// resets to verified completed bytes; cached files never trigger network access.
pub fn provision(directory: &Path, progress: &AtomicU64) -> Result<()> {
let manifest = manifest()?;
let client = Client::builder()
.connect_timeout(Duration::from_secs(20))
.timeout(Duration::from_secs(900))
.user_agent(concat!("Searchi/", env!("CARGO_PKG_VERSION")))
.build()?;
provision_with(directory, progress, &manifest, &client, |artifact| {
format!(
"https://huggingface.co/{}/resolve/{}/{}",
manifest.export_repository, manifest.revision, artifact.path
)
})
}
// URL injection is private so tests can use a local HTTP fixture. Production
// exposes no arbitrary download source or unverified executable provisioner.
fn provision_with(
directory: &Path,
progress: &AtomicU64,
manifest: &Manifest,
client: &Client,
url: impl Fn(&Artifact) -> String,
) -> Result<()> {
fs::create_dir_all(directory).context("create model cache directory")?;
progress.store(0, Ordering::Release);
let mut completed = 0;
for artifact in &manifest.files {
validate_artifact(artifact)?;
let destination = directory.join(&artifact.path);
if verified(&destination, artifact)? {
completed += artifact.bytes;
progress.store(completed, Ordering::Release);
continue;
}
let mut last_error = None;
for attempt in 0..ATTEMPTS {
match download_one(
&destination,
artifact,
progress,
completed,
client,
&url(artifact),
) {
Ok(()) => {
last_error = None;
break;
}
Err(error) => {
last_error = Some(error);
progress.store(completed, Ordering::Release);
if attempt + 1 < ATTEMPTS {
std::thread::sleep(Duration::from_millis(250 << attempt));
}
}
}
}
if let Some(error) = last_error {
return Err(error).with_context(|| {
format!(
"model artifact {} failed after {ATTEMPTS} attempts",
artifact.path
)
});
}
completed += artifact.bytes;
progress.store(completed, Ordering::Release);
}
Ok(())
}
fn download_one(
destination: &Path,
artifact: &Artifact,
progress: &AtomicU64,
completed: u64,
client: &Client,
url: &str,
) -> Result<()> {
let mut response = client
.get(url)
.send()
.and_then(reqwest::blocking::Response::error_for_status)
.map_err(reqwest::Error::without_url)
.context("request model artifact")?;
if response
.content_length()
.is_some_and(|size| size != artifact.bytes)
{
bail!("model artifact response length differs from pinned size");
}
let directory = destination
.parent()
.context("model artifact has no parent directory")?;
let mut output = tempfile::Builder::new()
.prefix(".searchi-model-")
.suffix(".part")
.tempfile_in(directory)?;
let mut hash = Sha256::new();
let mut downloaded = 0;
let mut buffer = [0u8; 65536];
loop {
let read = response
.read(&mut buffer)
.context("read model artifact stream")?;
if read == 0 {
break;
}
downloaded += read as u64;
if downloaded > artifact.bytes {
bail!("model artifact exceeds pinned size");
}
output.write_all(&buffer[..read])?;
hash.update(&buffer[..read]);
progress.store(completed + downloaded, Ordering::Release);
}
if downloaded != artifact.bytes
|| !format!("{:x}", hash.finalize()).eq_ignore_ascii_case(&artifact.sha256)
{
bail!("model artifact size or SHA-256 mismatch");
}
output.as_file().sync_all()?;
output
.persist(destination)
.map_err(|error| error.error)
.context("atomically replace model artifact")?;
#[cfg(unix)]
File::open(directory)?.sync_all()?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use std::{net::TcpListener, sync::atomic::AtomicUsize, thread, time::Instant};
fn artifact(path: &str, bytes: &[u8]) -> Artifact {
Artifact {
path: path.into(),
bytes: bytes.len() as u64,
sha256: format!("{:x}", Sha256::digest(bytes)),
}
}
fn test_manifest(files: Vec<Artifact>) -> Manifest {
Manifest {
model_id: "test".into(),
export_repository: "test".into(),
revision: "test".into(),
files,
}
}
fn client() -> Client {
Client::builder()
.no_proxy()
.connect_timeout(Duration::from_secs(1))
.timeout(Duration::from_secs(2))
.build()
.unwrap()
}
// Local HTTP only. The production API does not expose this override.
fn server(responses: Vec<Vec<u8>>) -> (String, thread::JoinHandle<usize>) {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let url = format!("http://{}/fixture", listener.local_addr().unwrap());
listener.set_nonblocking(true).unwrap();
let handle = thread::spawn(move || {
let mut count = 0;
for response in responses {
let deadline = Instant::now() + Duration::from_secs(5);
let mut stream = loop {
match listener.accept() {
Ok((stream, _)) => break stream,
Err(error)
if error.kind() == std::io::ErrorKind::WouldBlock
&& Instant::now() < deadline =>
{
thread::sleep(Duration::from_millis(2))
}
Err(error) => panic!("local fixture accept failed: {error}"),
}
};
stream.set_nonblocking(false).unwrap();
stream
.set_read_timeout(Some(Duration::from_secs(1)))
.unwrap();
let mut request = Vec::new();
let mut buffer = [0u8; 1024];
while !request.windows(4).any(|window| window == b"\r\n\r\n") {
let read = stream.read(&mut buffer).unwrap();
assert!(read > 0);
request.extend_from_slice(&buffer[..read]);
assert!(request.len() <= 8192);
}
stream.write_all(&response).unwrap();
count += 1;
}
count
});
(url, handle)
}
fn response(body: &[u8], declared: Option<usize>) -> Vec<u8> {
let mut response = b"HTTP/1.1 200 OK\r\nConnection: close\r\n".to_vec();
if let Some(length) = declared {
response.extend_from_slice(format!("Content-Length: {length}\r\n").as_bytes());
}
response.extend_from_slice(b"\r\n");
response.extend_from_slice(body);
response
}
fn assert_no_part_files(directory: &Path) {
assert!(!fs::read_dir(directory).unwrap().any(|entry| entry
.unwrap()
.file_name()
.to_string_lossy()
.ends_with(".part")));
}
#[test]
fn manifest_is_exact_pinned_google_text_and_media_contract() {
let manifest = manifest().unwrap();
assert_eq!(
manifest.revision,
"bfcd298762cc34d0357ece5ebdd31791a3a374d8"
);
assert_eq!(
manifest
.files
.iter()
.map(|artifact| artifact.bytes)
.sum::<u64>(),
TOTAL_BYTES
);
assert_eq!(
manifest.files[0].sha256,
"2188ac1deca4b77dffefd603c2776a9d76d9d74ec01841392982ebb840b09135"
);
assert_eq!(
manifest.files[1].sha256,
"c4a8a52691ecef40618438928bdf9e68379b854e24166f292592353db0aab64f"
);
}
#[test]
fn same_size_corruption_is_detected_and_missing_cache_is_offline() {
let directory = tempfile::tempdir().unwrap();
assert!(!verify_cached(directory.path()).unwrap());
let artifact = artifact("model.gguf", b"abc");
let path = directory.path().join(&artifact.path);
fs::write(&path, b"abc").unwrap();
assert!(verified(&path, &artifact).unwrap());
fs::write(&path, b"xyz").unwrap();
assert!(!verified(&path, &artifact).unwrap());
}
#[test]
fn cached_valid_artifacts_do_not_request_a_url() {
let directory = tempfile::tempdir().unwrap();
fs::write(directory.path().join("one.gguf"), b"abc").unwrap();
let manifest = test_manifest(vec![artifact("one.gguf", b"abc")]);
let calls = AtomicUsize::new(0);
let progress = AtomicU64::new(999);
provision_with(directory.path(), &progress, &manifest, &client(), |_| {
calls.fetch_add(1, Ordering::Relaxed);
"http://127.0.0.1:1/unreachable".into()
})
.unwrap();
assert_eq!(calls.load(Ordering::Relaxed), 0);
assert_eq!(progress.load(Ordering::Acquire), 3);
}
#[test]
fn interrupted_attempt_retries_and_atomically_replaces_corrupt_cache() {
let directory = tempfile::tempdir().unwrap();
fs::write(directory.path().join("one.gguf"), b"badbad").unwrap();
fs::write(directory.path().join("two.gguf"), b"abcdef").unwrap();
let manifest = test_manifest(vec![
artifact("one.gguf", b"abcdef"),
artifact("two.gguf", b"abcdef"),
]);
let (url, server) = server(vec![
response(b"abc", Some(6)),
response(b"abcdef", Some(6)),
]);
let progress = AtomicU64::new(0);
provision_with(directory.path(), &progress, &manifest, &client(), |_| {
url.clone()
})
.unwrap();
assert_eq!(server.join().unwrap(), 2);
assert_eq!(progress.load(Ordering::Acquire), 12);
assert_eq!(
fs::read(directory.path().join("one.gguf")).unwrap(),
b"abcdef"
);
assert_no_part_files(directory.path());
}
#[test]
fn oversized_stream_is_bounded_and_failure_cleans_partial_files() {
let directory = tempfile::tempdir().unwrap();
fs::write(directory.path().join("one.gguf"), b"old").unwrap();
let manifest = test_manifest(vec![artifact("one.gguf", b"abc")]);
let (url, server) = server(vec![response(b"abcde", None); ATTEMPTS]);
let progress = AtomicU64::new(0);
let error = provision_with(directory.path(), &progress, &manifest, &client(), |_| {
url.clone()
})
.unwrap_err();
assert!(format!("{error:#}").contains("exceeds pinned size"));
assert_eq!(server.join().unwrap(), ATTEMPTS);
assert_eq!(progress.load(Ordering::Acquire), 0);
assert_eq!(fs::read(directory.path().join("one.gguf")).unwrap(), b"old");
assert_no_part_files(directory.path());
}
#[test]
fn wrong_digest_cannot_be_installed_even_with_correct_size() {
let directory = tempfile::tempdir().unwrap();
let manifest = test_manifest(vec![artifact("one.gguf", b"abc")]);
let (url, server) = server(vec![response(b"xyz", Some(3)); ATTEMPTS]);
let progress = AtomicU64::new(0);
let error = provision_with(directory.path(), &progress, &manifest, &client(), |_| {
url.clone()
})
.unwrap_err();
assert!(format!("{error:#}").contains("SHA-256 mismatch"));
assert_eq!(server.join().unwrap(), ATTEMPTS);
assert!(!directory.path().join("one.gguf").exists());
assert_no_part_files(directory.path());
}
#[test]
fn completed_artifact_survives_later_failure_and_is_reused_on_retry() {
let directory = tempfile::tempdir().unwrap();
let manifest = test_manifest(vec![
artifact("one.gguf", b"abc"),
artifact("two.gguf", b"def"),
]);
let mut responses = vec![response(b"abc", Some(3))];
responses.extend(vec![response(b"bad", Some(3)); ATTEMPTS]);
let (url, first_server) = server(responses);
let progress = AtomicU64::new(0);
assert!(
provision_with(directory.path(), &progress, &manifest, &client(), |_| url
.clone())
.is_err()
);
assert_eq!(first_server.join().unwrap(), ATTEMPTS + 1);
assert_eq!(progress.load(Ordering::Acquire), 3);
assert_eq!(fs::read(directory.path().join("one.gguf")).unwrap(), b"abc");
assert!(!directory.path().join("two.gguf").exists());
let (url, second_server) = server(vec![response(b"def", Some(3))]);
provision_with(
directory.path(),
&progress,
&manifest,
&client(),
|artifact| {
assert_eq!(artifact.path, "two.gguf");
url.clone()
},
)
.unwrap();
assert_eq!(second_server.join().unwrap(), 1);
assert_eq!(progress.load(Ordering::Acquire), 6);
assert_no_part_files(directory.path());
}
#[test]
fn unsafe_artifact_names_are_rejected_and_default_path_is_platform_local() {
assert!(validate_artifact(&artifact("../outside", b"abc")).is_err());
assert!(validate_artifact(&artifact("..\\outside", b"abc")).is_err());
assert!(default_model_directory()
.unwrap()
.ends_with("models/embeddinggemma2"));
}
}
crates/searchi-cli/Cargo.toml
+2
View File
@@ -13,4 +13,6 @@ path = "src/main.rs"
anyhow.workspace = true
clap = { version = "4", features = ["derive"] }
searchi-core = { path = "../searchi-core" }
searchi-service = { path = "../searchi-service" }
searchi-assets = { path = "../searchi-assets" }
serde_json.workspace = true
crates/searchi-cli/src/main.rs
+96 -16
View File
@@ -5,6 +5,7 @@ use std::path::PathBuf;
use anyhow::{bail, Context, Result};
use clap::{Parser, Subcommand};
use searchi_core::{IndexOptions, SearchIndex, SearchOptions};
use searchi_service::HybridSearch;
#[derive(Debug, Parser)]
#[command(name = "searchi", version, about = "Search your local files offline")]
@@ -15,6 +16,12 @@ struct Cli {
#[derive(Debug, Subcommand)]
enum Command {
/// Set up the required Google model in the default local cache.
Setup {
/// 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.
@@ -54,6 +61,25 @@ enum Command {
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,
},
/// Show the indexed file and content counts.
Stats {
@@ -66,6 +92,15 @@ enum Command {
fn main() -> Result<()> {
match Cli::parse().command {
Command::Setup { model_dir } => {
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);
@@ -74,6 +109,9 @@ fn main() -> Result<()> {
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
@@ -90,27 +128,32 @@ fn main() -> Result<()> {
max_size,
modified_after_ms,
modified_before_ms,
model_dir,
embedding_worker,
cpu_threads,
} => {
if query.trim().is_empty() {
bail!("Search text cannot be empty");
}
require_existing_index(&db)?;
let index = SearchIndex::open(&db)
.with_context(|| format!("Cannot open index {}", db.display()))?;
let results = index
.search_with_options(
&query,
usize::from(limit),
&SearchOptions {
folder,
extension,
min_size,
max_size,
modified_after_unix_ms: modified_after_ms,
modified_before_unix_ms: modified_before_ms,
},
)
.context("Search failed")?;
let options = SearchOptions {
folder,
extension,
min_size,
max_size,
modified_after_unix_ms: modified_after_ms,
modified_before_unix_ms: modified_before_ms,
};
let results = if let (Some(model), Some(runtime)) = (model_dir, embedding_worker) {
HybridSearch::open(&db, model, runtime, usize::from(cpu_threads))?
.search(&query, usize::from(limit), &options)
.context("Hybrid search failed")?
} else {
SearchIndex::open(&db)
.with_context(|| format!("Cannot open index {}", db.display()))?
.search_with_options(&query, usize::from(limit), &options)
.context("Search failed")?
};
if json {
println!("{}", serde_json::to_string_pretty(&results)?);
} else if results.is_empty() {
@@ -130,6 +173,19 @@ fn main() -> Result<()> {
}
}
}
Command::Embed {
db,
model_dir,
embedding_worker,
cpu_threads,
} => {
require_existing_index(&db)?;
let report =
HybridSearch::open(&db, model_dir, embedding_worker, usize::from(cpu_threads))?
.index_embeddings()
.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)
@@ -173,6 +229,30 @@ mod tests {
}
}
#[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([
crates/searchi-core/README.md
+24 -10
View File
@@ -4,7 +4,7 @@ The core has no desktop or inference dependencies. Its current retrieval engine
combines literal filename matching with SQLite FTS5 BM25 keyword search. It does
not implement semantic search.
`SearchIndex::open` creates or migrates schema version 7. Current-schema opens do
`SearchIndex::open` creates or migrates schema version 9. Current-schema opens do
not acquire a write lock. WAL allows independent readers to use the committed
index while reconciliation runs. All indexing changes are committed atomically;
a traversal or plain-text file read error rolls back the entire reconciliation.
@@ -82,20 +82,33 @@ Exact filenames always precede other lanes. Ties use persistent file row IDs,
and keyword ties use relevance then row ID. Selected rows are deduplicated.
Filename/path substrings of at least three Unicode characters use an
external-content trigram FTS index. The rarest known query trigram supplies a
candidate superset; the whole literal substring is verified on the candidate.
external-content trigram FTS index. A positional anchor grows around the rarest
known query trigram, admitting only neighboring grams within four times its
document frequency (with a small floor) and at most 32 Unicode characters.
This rejects repeated-identifier false candidates inside FTS without extending
through dense common suffix postings. The anchor supplies a candidate superset;
the whole literal substring is still verified on each candidate.
At most 256 unique query trigrams are inspected. This lookup budget does not cap
candidate rows or hide matching results. All size/type/folder/date filters apply
before each lane's SQL LIMIT. Shorter substrings retain complete metadata scans
and can be slow. Broad queries under restrictive filters can also examine many
matching candidates; no universal filtered-query latency claim is made.
Trigram frequencies use the FTS vocabulary and a bounded 4,096-term cache on each
`SearchIndex`. Initial frequency lookups can be expensive; the cache is cleared
after reconciliation. Keep an independent reader engine alive across searches to
benefit from warm behavior. These frequency estimates affect candidate selection
cost, not correctness. A source query batches all finalists' indexed passage IDs
and selects actual matching passages before choosing one source per file.
`prewarm()` refreshes persisted trigram frequencies from the FTS vocabulary only
when names/path membership is dirty, then loads a bounded 4,096-term reader cache.
Foreground frequency misses use an indexed persisted table, not vocabulary
postings scans. Name/path insert/delete/change marks statistics dirty; content
changes alone do not. Statistics affect candidate cost, not match correctness.
Background preparation exercises actual first/last filenames, alphabetic name
tokens, and three words of the first document. UI callers should complete this
preparation before enabling interaction. Representative inputs are disclosed in
the benchmark report; preparation is not counted as foreground query time.
Unfiltered content retrieval materializes ranked row IDs before reading finalist
metadata and snippets. Filtered requests retain complete filter-before-LIMIT
ranking. Tests compare ordering, scores, snippets, and sources between both paths.
A source query batches all finalists' indexed passage IDs and chooses one actual
matching source per file. Keep an independent reader alive across searches.
Migration to the new indexes preserves file content and passage IDs and rebuilds
both filename/path and passage indexes atomically. Large migrations can take time
@@ -105,7 +118,8 @@ opens. The measured Windows synthetic fixture grew from about 295 MiB to 605 MiB
## Known limits
- Short substring queries and broad restrictive filters can still scan many
metadata candidates. Cold frequency-cache fills exceed the warm filename target.
metadata candidates. Statistics preparation may take seconds after large scans;
arbitrary unseen or OS-cold queries have no universal latency guarantee.
- Files changed without changing size or modification time require a future
forced refresh/content-fingerprint option. Timestamp precision depends on the
filesystem. Indexing does not yet run a filesystem watcher.
crates/searchi-core/benchmarks/README.md
+124
View File
@@ -102,3 +102,127 @@ is still not established. Query distribution and filters are not represented by
this four-case fixture; restrictive broad filters and sub-three-character searches
remain potentially slow, while correctness tests ensure late matching files are
not lost through arbitrary pre-filter caps.
## Actual first foreground inputs, schema 8
The schema-8 run persists trigram frequency statistics in background preparation
before allowing the measured inputs. Preparation took **17.21 seconds** and
loaded 1,268 frequency entries. That time is reported separately and is not
subtracted from a foreground query. The same 1M/100K fixture was reused.
The harness measures 20 different exact filenames, spread across the corpus,
with no preceding query warmups. It then measures 20 different filenames, each
on a newly opened reader without calling `prewarm` again, using persisted stats.
This second set includes reader-open time. Both sets validate exact-match priority
and record every raw query and duration. They are disjoint from each other and
from the four-case exact query for this full fixture.
| Distinct first-input set | p50 | p95 | Actual very first input | Samples ≥50 ms |
| --- | ---: | ---: | ---: | ---: |
| Prepared reader, 20 names | 13.96 ms | 18.97 ms | **98.49 ms** | 1/20 |
| Fresh readers, 20 other names | 16.78 ms | 22.73 ms | **103.57 ms** | 1/20 |
The p95 filename targets pass, but the very first zero-heavy filename in each
set fails the 50 ms target. The nearest-rank p95 excludes one maximum sample
when there are 20 samples, so p95 alone must not be used to claim that the actual
first input is always fast. Explicit first-sample flags and maximum durations
are included in the report, derived directly from the recorded raw samples.
| Existing case, first execution before its warmups | First input | First target | Repeated warm p95 |
| --- | ---: | --- | ---: |
| Exact filename at end of corpus | 16.31 ms | Passed <50 ms | 13.80 ms |
| Broad filename substring | 13.01 ms | Passed <50 ms | 0.47 ms |
| Unique content token | 14.07 ms | Passed <150 ms | 0.49 ms |
| Broad keyword | **264.53 ms** | **Failed <150 ms** | 145.19 ms |
These measurements expose remaining first-input work even after frequency
preparation. Zero-heavy names produce larger candidate sets, and broad content
retrieval has a first-use penalty. This report does not establish the requested
fast-first-input behavior as complete. It also does not represent a cold OS
page-cache test: background preparation and previous operations can warm disk
pages. Host workload isolation remains unverified. The retained report is
`windows-cpu-full-prepared-v8.json`; engine code was not changed during this run.
## Background representative search and restart cost, schema 9
Schema 9 refreshes frequency statistics only when filename/path membership is
dirty. It also runs representative searches in background. The report explicitly
records the actual five background inputs from this corpus: `invoice_0000000.txt`,
`invoice_0999999.txt`, `Supplier`, `termination`, and `cancellation`. Consequently
the later first foreground executions of overlapping inputs are prepared inputs,
not never-before-executed queries. This behavior is intentional application
preparation, rather than an undisclosed benchmark warmup.
Initial preparation took **17.723 seconds**. A second `prewarm()` on a newly
opened reader took **0.396 seconds**, reusing clean persisted statistics. That
second call occurred only after all foreground measurements, so it could not
improve their recorded first timings. Startup time is amortized when names have
not changed; refresh after mutations still has an explicit background cost.
| Twenty distinct foreground names | p50 | p95 | Very first name | ≥50 ms |
| --- | ---: | ---: | ---: | ---: |
| Prepared reader | 12.75 ms | 14.40 ms | **85.38 ms** | 1/20 |
| Newly reopened readers | 16.31 ms | 31.42 ms | **100.07 ms** | 1/20 |
The remaining zero-heavy filename outliers fail the first-input target even
after representative preparation. Repeating that representative query in
background does not eliminate its candidate-processing cost. The 20 distinct
query sets remain unchanged from schema 8, and every raw sample and explicit
first/max failure flag is retained.
| Existing case | First foreground | Repeated p95 | Target |
| --- | ---: | ---: | --- |
| Exact filename at corpus end | 13.48 ms | 13.22 ms | Passed <50 ms |
| Filename substring | 3.20 ms | 0.43 ms | Passed <50 ms |
| Unique content token | 2.57 ms | 0.53 ms | Passed <150 ms |
| Broad keyword | 99.10 ms | 106.20 ms | Passed <150 ms |
Background initialization fixes the observed broad-keyword first-use penalty,
but the full fast-first-input requirement remains incomplete because of the
filename outliers. Further algorithmic filename work is required; the report
does not infer that preparation makes every unseen query fast. This Windows
CPU-only run is stored in `windows-cpu-full-prepared-v9.json`.
## Positional candidate correction and finalist projection
Extending the rarest trigram through the entire suffix regressed retrieval:
the unchanged 20-name sets had p95 **136.60/129.17 ms**, and the first reopened
zero-heavy name took **554.32 ms**. That rejected experiment remains recorded in
`windows-cpu-full-suffix-anchor.json`. A long phrase can initialize dense suffix
postings even when its first term is rare.
The replacement grows a positional anchor only through neighboring grams within
a bounded frequency factor, with a maximum of 32 Unicode characters. Full literal
verification and filter-before-limit completeness remain unchanged. Temporary
lane profiling separated cheap indexed equality (about 0.06 ms in the diagnostic
run) from first-use broad-prefix filename retrieval (about 67 ms). Generic
alphabetic tokens derived from representative names now initialize those common
postings in background. Unfiltered content searches rank compact IDs before
materializing finalist paths/snippets; filtered ranking is kept complete.
The final unchanged 1M/100K corpus run uses these disclosed background inputs:
first/last filenames, `invoice`, `txt`, and the first document words `Supplier`,
`termination`, and `cancellation`. Clean-statistics preparation took **0.709 s**;
a subsequent new-reader preparation took **0.264 s**. Earlier initial dirty
statistics refresh still took roughly 17 seconds; this run does not erase that
startup/refresh cost.
| Twenty distinct first inputs | p95 | Very first | Maximum | Samples ≥50 ms |
| --- | ---: | ---: | ---: | ---: |
| Prepared reader | 18.15 ms | 6.14 ms | 22.94 ms | 0/20 |
| Newly reopened readers, including open | 30.86 ms | 30.86 ms | 34.28 ms | 0/20 |
| Existing case | First foreground | Repeated p95 | Target |
| --- | ---: | ---: | --- |
| Exact filename at corpus end | 4.37 ms | 3.77 ms | Passed <50 ms |
| Filename substring | 14.18 ms | 1.23 ms | Passed <50 ms |
| Unique content token | 10.06 ms | 0.72 ms | Passed <150 ms |
| Broad keyword | 114.44 ms | 117.88 ms | Passed <150 ms |
All measured first-input samples and repeated-query p95 values pass in this
Windows CPU run. The retained final report is
`windows-cpu-full-candidate-first-final.json`, with every raw filename/sample.
This establishes the measured fixture result, not a guarantee for all queries,
filters, corpus distributions, cold OS caches, or other operating systems. The
intermediate density-only reports retain remaining failures and are not promoted
as final acceptance evidence.
crates/searchi-core/benchmarks/results/windows-cpu-full-candidate-first-final.json
+179
View File
@@ -0,0 +1,179 @@
{
"background_prepare_seconds": 0.7093476,
"reopened_background_prepare_seconds": 0.2639451,
"background_representative_queries": [
"invoice_0000000.txt",
"invoice_0999999.txt",
"invoice",
"txt",
"Supplier",
"termination",
"cancellation"
],
"prepared_frequency_count": 1268,
"first_foreground_distinct_names": {
"description": "distinct foreground filenames; no benchmark query warmups; production background representative queries may overlap",
"includes_reader_open": false,
"queries": [
"invoice_0000000.txt",
"invoice_0049999.txt",
"invoice_0099999.txt",
"invoice_0149999.txt",
"invoice_0199999.txt",
"invoice_0249999.txt",
"invoice_0299999.txt",
"invoice_0349999.txt",
"invoice_0399999.txt",
"invoice_0449999.txt",
"invoice_0499999.txt",
"invoice_0549998.txt",
"invoice_0599998.txt",
"invoice_0649998.txt",
"invoice_0699998.txt",
"invoice_0749998.txt",
"invoice_0799998.txt",
"invoice_0849998.txt",
"invoice_0899998.txt",
"invoice_0949998.txt"
],
"samples_ms": [
6.1442,
18.1445,
3.5369,
14.9677,
12.8753,
13.4905,
9.2402,
10.4807,
8.5722,
11.1046,
4.2845,
22.9355,
12.088099999999999,
13.8535,
9.436399999999999,
13.2516,
11.2014,
10.591899999999999,
9.587399999999999,
9.6533
],
"p50_ms": 10.591899999999999,
"p95_ms": 18.1445,
"first_sample_ms": 6.1442,
"first_sample_target_met": true,
"maximum_ms": 22.9355,
"samples_exceeding_target": 0,
"target_p95_ms": 50.0,
"target_met": true
},
"reopened_reader_first_names": {
"description": "each never-repeated filename on a newly opened reader, persisted statistics, no prewarm; includes open+search",
"includes_reader_open": true,
"queries": [
"invoice_0000001.txt",
"invoice_0050000.txt",
"invoice_0100000.txt",
"invoice_0150000.txt",
"invoice_0200000.txt",
"invoice_0250000.txt",
"invoice_0300000.txt",
"invoice_0350000.txt",
"invoice_0400000.txt",
"invoice_0450000.txt",
"invoice_0500000.txt",
"invoice_0549999.txt",
"invoice_0599999.txt",
"invoice_0649999.txt",
"invoice_0699999.txt",
"invoice_0749999.txt",
"invoice_0799999.txt",
"invoice_0849999.txt",
"invoice_0899999.txt",
"invoice_0949999.txt"
],
"samples_ms": [
30.8552,
28.1971,
34.2776,
14.4363,
29.89,
13.669,
28.3825,
9.8496,
28.3787,
13.243400000000001,
21.4598,
7.0038,
6.1826,
6.6606,
6.697,
6.902,
7.950900000000001,
6.4144,
6.8307,
7.2419
],
"p50_ms": 9.8496,
"p95_ms": 30.8552,
"first_sample_ms": 30.8552,
"first_sample_target_met": true,
"maximum_ms": 34.2776,
"samples_exceeding_target": 0,
"target_p95_ms": 50.0,
"target_met": true
},
"fixture": "synthetic SQLite metadata and text; no filesystem extraction or semantic inference",
"files": 1000000,
"documents": 100000,
"seed_seconds": 0.0,
"migration_seconds": 0.0110352,
"fixture_reused": true,
"database_bytes": 634449920,
"database_path": "target/benchmarks/full.db",
"architecture": "x86_64",
"os": "windows",
"cpu_only": true,
"queries": [
{
"query": "invoice_0999999.txt",
"samples": 20,
"p50_ms": 2.3605,
"first_query_ms": 4.368399999999999,
"first_query_target_met": true,
"p95_ms": 3.7694,
"target_p95_ms": 50.0,
"target_met": true
},
{
"query": "invoice_",
"samples": 20,
"p50_ms": 0.7753,
"first_query_ms": 14.1849,
"first_query_target_met": true,
"p95_ms": 1.228,
"target_p95_ms": 50.0,
"target_met": true
},
{
"query": "marker0099999",
"samples": 20,
"p50_ms": 0.47550000000000003,
"first_query_ms": 10.0635,
"first_query_target_met": true,
"p95_ms": 0.7228,
"target_p95_ms": 150.0,
"target_met": true
},
{
"query": "cancellation",
"samples": 20,
"p50_ms": 101.68339999999999,
"first_query_ms": 114.44120000000001,
"first_query_target_met": true,
"p95_ms": 117.8776,
"target_p95_ms": 150.0,
"target_met": true
}
]
}
crates/searchi-core/benchmarks/results/windows-cpu-full-density-anchor.json
+177
View File
@@ -0,0 +1,177 @@
{
"background_prepare_seconds": 0.7323155,
"reopened_background_prepare_seconds": 0.3083516,
"background_representative_queries": [
"invoice_0000000.txt",
"invoice_0999999.txt",
"Supplier",
"termination",
"cancellation"
],
"prepared_frequency_count": 1268,
"first_foreground_distinct_names": {
"description": "distinct foreground filenames; no benchmark query warmups; production background representative queries may overlap",
"includes_reader_open": false,
"queries": [
"invoice_0000000.txt",
"invoice_0049999.txt",
"invoice_0099999.txt",
"invoice_0149999.txt",
"invoice_0199999.txt",
"invoice_0249999.txt",
"invoice_0299999.txt",
"invoice_0349999.txt",
"invoice_0399999.txt",
"invoice_0449999.txt",
"invoice_0499999.txt",
"invoice_0549998.txt",
"invoice_0599998.txt",
"invoice_0649998.txt",
"invoice_0699998.txt",
"invoice_0749998.txt",
"invoice_0799998.txt",
"invoice_0849998.txt",
"invoice_0899998.txt",
"invoice_0949998.txt"
],
"samples_ms": [
5.4780999999999995,
14.290199999999999,
2.2094,
9.4278,
9.712800000000001,
11.698,
7.796,
9.2627,
7.9072,
71.67869999999999,
2.4205,
20.3763,
9.2537,
11.2855,
7.5798000000000005,
9.6667,
7.8527000000000005,
8.970600000000001,
5.9854,
8.1966
],
"p50_ms": 8.970600000000001,
"p95_ms": 20.3763,
"first_sample_ms": 5.4780999999999995,
"first_sample_target_met": true,
"maximum_ms": 71.67869999999999,
"samples_exceeding_target": 1,
"target_p95_ms": 50.0,
"target_met": true
},
"reopened_reader_first_names": {
"description": "each never-repeated filename on a newly opened reader, persisted statistics, no prewarm; includes open+search",
"includes_reader_open": true,
"queries": [
"invoice_0000001.txt",
"invoice_0050000.txt",
"invoice_0100000.txt",
"invoice_0150000.txt",
"invoice_0200000.txt",
"invoice_0250000.txt",
"invoice_0300000.txt",
"invoice_0350000.txt",
"invoice_0400000.txt",
"invoice_0450000.txt",
"invoice_0500000.txt",
"invoice_0549999.txt",
"invoice_0599999.txt",
"invoice_0649999.txt",
"invoice_0699999.txt",
"invoice_0749999.txt",
"invoice_0799999.txt",
"invoice_0849999.txt",
"invoice_0899999.txt",
"invoice_0949999.txt"
],
"samples_ms": [
12.800699999999999,
26.5927,
25.4506,
11.3757,
24.2376,
10.9734,
22.7002,
8.4673,
22.9176,
11.1351,
16.4276,
5.4818999999999996,
4.0053,
4.1241,
3.9019,
4.209,
3.9389000000000003,
4.0844000000000005,
3.8915,
5.3457
],
"p50_ms": 8.4673,
"p95_ms": 25.4506,
"first_sample_ms": 12.800699999999999,
"first_sample_target_met": true,
"maximum_ms": 26.5927,
"samples_exceeding_target": 0,
"target_p95_ms": 50.0,
"target_met": true
},
"fixture": "synthetic SQLite metadata and text; no filesystem extraction or semantic inference",
"files": 1000000,
"documents": 100000,
"seed_seconds": 0.0,
"migration_seconds": 0.0040623,
"fixture_reused": true,
"database_bytes": 634449920,
"database_path": "target/benchmarks/full.db",
"architecture": "x86_64",
"os": "windows",
"cpu_only": true,
"queries": [
{
"query": "invoice_0999999.txt",
"samples": 20,
"p50_ms": 1.3525,
"first_query_ms": 2.9366000000000003,
"first_query_target_met": true,
"p95_ms": 2.5606,
"target_p95_ms": 50.0,
"target_met": true
},
{
"query": "invoice_",
"samples": 20,
"p50_ms": 0.5368999999999999,
"first_query_ms": 66.6776,
"first_query_target_met": false,
"p95_ms": 0.7038,
"target_p95_ms": 50.0,
"target_met": true
},
{
"query": "marker0099999",
"samples": 20,
"p50_ms": 0.2814,
"first_query_ms": 10.037799999999999,
"first_query_target_met": true,
"p95_ms": 0.4111,
"target_p95_ms": 150.0,
"target_met": true
},
{
"query": "cancellation",
"samples": 20,
"p50_ms": 98.19590000000001,
"first_query_ms": 107.1098,
"first_query_target_met": true,
"p95_ms": 102.2884,
"target_p95_ms": 150.0,
"target_met": true
}
]
}
crates/searchi-core/benchmarks/results/windows-cpu-full-density-prepared.json
+179
View File
@@ -0,0 +1,179 @@
{
"background_prepare_seconds": 0.9157306,
"reopened_background_prepare_seconds": 0.390985,
"background_representative_queries": [
"invoice_0000000.txt",
"invoice_0999999.txt",
"invoice",
"txt",
"Supplier",
"termination",
"cancellation"
],
"prepared_frequency_count": 1268,
"first_foreground_distinct_names": {
"description": "distinct foreground filenames; no benchmark query warmups; production background representative queries may overlap",
"includes_reader_open": false,
"queries": [
"invoice_0000000.txt",
"invoice_0049999.txt",
"invoice_0099999.txt",
"invoice_0149999.txt",
"invoice_0199999.txt",
"invoice_0249999.txt",
"invoice_0299999.txt",
"invoice_0349999.txt",
"invoice_0399999.txt",
"invoice_0449999.txt",
"invoice_0499999.txt",
"invoice_0549998.txt",
"invoice_0599998.txt",
"invoice_0649998.txt",
"invoice_0699998.txt",
"invoice_0749998.txt",
"invoice_0799998.txt",
"invoice_0849998.txt",
"invoice_0899998.txt",
"invoice_0949998.txt"
],
"samples_ms": [
6.9921999999999995,
16.238,
2.3243,
13.4229,
9.434,
11.145,
8.304300000000001,
11.0419,
9.932699999999999,
10.8164,
2.3062,
21.6846,
9.497,
19.0655,
8.6778,
14.024099999999999,
10.106300000000001,
11.553799999999999,
8.5308,
7.7924999999999995
],
"p50_ms": 9.932699999999999,
"p95_ms": 19.0655,
"first_sample_ms": 6.9921999999999995,
"first_sample_target_met": true,
"maximum_ms": 21.6846,
"samples_exceeding_target": 0,
"target_p95_ms": 50.0,
"target_met": true
},
"reopened_reader_first_names": {
"description": "each never-repeated filename on a newly opened reader, persisted statistics, no prewarm; includes open+search",
"includes_reader_open": true,
"queries": [
"invoice_0000001.txt",
"invoice_0050000.txt",
"invoice_0100000.txt",
"invoice_0150000.txt",
"invoice_0200000.txt",
"invoice_0250000.txt",
"invoice_0300000.txt",
"invoice_0350000.txt",
"invoice_0400000.txt",
"invoice_0450000.txt",
"invoice_0500000.txt",
"invoice_0549999.txt",
"invoice_0599999.txt",
"invoice_0649999.txt",
"invoice_0699999.txt",
"invoice_0749999.txt",
"invoice_0799999.txt",
"invoice_0849999.txt",
"invoice_0899999.txt",
"invoice_0949999.txt"
],
"samples_ms": [
27.8174,
25.6445,
29.525699999999997,
11.5664,
27.617,
10.4947,
30.0687,
10.766,
28.287300000000002,
13.547,
23.8423,
6.6371,
4.4586,
4.5778,
5.1159,
4.6030999999999995,
4.845000000000001,
5.5757,
4.6811,
6.2597
],
"p50_ms": 10.4947,
"p95_ms": 29.525699999999997,
"first_sample_ms": 27.8174,
"first_sample_target_met": true,
"maximum_ms": 30.0687,
"samples_exceeding_target": 0,
"target_p95_ms": 50.0,
"target_met": true
},
"fixture": "synthetic SQLite metadata and text; no filesystem extraction or semantic inference",
"files": 1000000,
"documents": 100000,
"seed_seconds": 0.0,
"migration_seconds": 0.0098789,
"fixture_reused": true,
"database_bytes": 634449920,
"database_path": "target/benchmarks/full.db",
"architecture": "x86_64",
"os": "windows",
"cpu_only": true,
"queries": [
{
"query": "invoice_0999999.txt",
"samples": 20,
"p50_ms": 1.8600999999999999,
"first_query_ms": 2.1075,
"first_query_target_met": true,
"p95_ms": 2.5455,
"target_p95_ms": 50.0,
"target_met": true
},
{
"query": "invoice_",
"samples": 20,
"p50_ms": 0.5207999999999999,
"first_query_ms": 12.9272,
"first_query_target_met": true,
"p95_ms": 0.6883,
"target_p95_ms": 50.0,
"target_met": true
},
{
"query": "marker0099999",
"samples": 20,
"p50_ms": 0.33340000000000003,
"first_query_ms": 10.2134,
"first_query_target_met": true,
"p95_ms": 0.4555,
"target_p95_ms": 150.0,
"target_met": true
},
{
"query": "cancellation",
"samples": 20,
"p50_ms": 141.4495,
"first_query_ms": 125.0308,
"first_query_target_met": true,
"p95_ms": 172.7458,
"target_p95_ms": 150.0,
"target_met": false
}
]
}
crates/searchi-core/benchmarks/results/windows-cpu-full-prepared-v8.json
+169
View File
@@ -0,0 +1,169 @@
{
"background_prepare_seconds": 17.2134568,
"prepared_frequency_count": 1268,
"first_foreground_distinct_names": {
"description": "never-repeated filenames on prepared reader; no query warmups",
"includes_reader_open": false,
"queries": [
"invoice_0000000.txt",
"invoice_0049999.txt",
"invoice_0099999.txt",
"invoice_0149999.txt",
"invoice_0199999.txt",
"invoice_0249999.txt",
"invoice_0299999.txt",
"invoice_0349999.txt",
"invoice_0399999.txt",
"invoice_0449999.txt",
"invoice_0499999.txt",
"invoice_0549998.txt",
"invoice_0599998.txt",
"invoice_0649998.txt",
"invoice_0699998.txt",
"invoice_0749998.txt",
"invoice_0799998.txt",
"invoice_0849998.txt",
"invoice_0899998.txt",
"invoice_0949998.txt"
],
"samples_ms": [
98.4897,
17.2344,
15.2996,
14.8814,
13.0075,
12.6754,
12.543099999999999,
13.686200000000001,
13.9632,
13.702300000000001,
14.8554,
15.3544,
15.874699999999999,
13.2766,
17.0515,
18.9665,
12.6579,
14.1467,
12.6395,
13.706100000000001
],
"p50_ms": 13.9632,
"p95_ms": 18.9665,
"first_sample_ms": 98.4897,
"first_sample_target_met": false,
"maximum_ms": 98.4897,
"samples_exceeding_target": 1,
"target_p95_ms": 50.0,
"target_met": true
},
"reopened_reader_first_names": {
"description": "each never-repeated filename on a newly opened reader, persisted statistics, no prewarm; includes open+search",
"includes_reader_open": true,
"queries": [
"invoice_0000001.txt",
"invoice_0050000.txt",
"invoice_0100000.txt",
"invoice_0150000.txt",
"invoice_0200000.txt",
"invoice_0250000.txt",
"invoice_0300000.txt",
"invoice_0350000.txt",
"invoice_0400000.txt",
"invoice_0450000.txt",
"invoice_0500000.txt",
"invoice_0549999.txt",
"invoice_0599999.txt",
"invoice_0649999.txt",
"invoice_0699999.txt",
"invoice_0749999.txt",
"invoice_0799999.txt",
"invoice_0849999.txt",
"invoice_0899999.txt",
"invoice_0949999.txt"
],
"samples_ms": [
103.57130000000001,
19.669,
17.804299999999998,
16.7811,
15.733600000000001,
16.381300000000003,
16.7549,
16.395500000000002,
21.7466,
17.044,
21.080399999999997,
22.7327,
20.9134,
18.464399999999998,
17.691300000000002,
14.9367,
15.176,
16.0537,
14.7705,
15.5895
],
"p50_ms": 16.7811,
"p95_ms": 22.7327,
"first_sample_ms": 103.57130000000001,
"first_sample_target_met": false,
"maximum_ms": 103.57130000000001,
"samples_exceeding_target": 1,
"target_p95_ms": 50.0,
"target_met": true
},
"fixture": "synthetic SQLite metadata and text; no filesystem extraction or semantic inference",
"files": 1000000,
"documents": 100000,
"seed_seconds": 0.0,
"migration_seconds": 0.0128952,
"fixture_reused": true,
"database_bytes": 634449920,
"database_path": "target/benchmarks/full.db",
"architecture": "x86_64",
"os": "windows",
"cpu_only": true,
"queries": [
{
"query": "invoice_0999999.txt",
"samples": 20,
"p50_ms": 12.3028,
"first_query_ms": 16.3077,
"first_query_target_met": true,
"p95_ms": 13.7969,
"target_p95_ms": 50.0,
"target_met": true
},
{
"query": "invoice_",
"samples": 20,
"p50_ms": 0.33149999999999996,
"first_query_ms": 13.0136,
"first_query_target_met": true,
"p95_ms": 0.4704,
"target_p95_ms": 50.0,
"target_met": true
},
{
"query": "marker0099999",
"samples": 20,
"p50_ms": 0.2775,
"first_query_ms": 14.0693,
"first_query_target_met": true,
"p95_ms": 0.49319999999999997,
"target_p95_ms": 150.0,
"target_met": true
},
{
"query": "cancellation",
"samples": 20,
"p50_ms": 101.3439,
"first_query_ms": 264.5315,
"first_query_target_met": false,
"p95_ms": 145.1926,
"target_p95_ms": 150.0,
"target_met": true
}
]
}
crates/searchi-core/benchmarks/results/windows-cpu-full-prepared-v9.json
+177
View File
@@ -0,0 +1,177 @@
{
"background_prepare_seconds": 17.7229177,
"reopened_background_prepare_seconds": 0.3963525,
"background_representative_queries": [
"invoice_0000000.txt",
"invoice_0999999.txt",
"Supplier",
"termination",
"cancellation"
],
"prepared_frequency_count": 1268,
"first_foreground_distinct_names": {
"description": "distinct foreground filenames; no benchmark query warmups; production background representative queries may overlap",
"includes_reader_open": false,
"queries": [
"invoice_0000000.txt",
"invoice_0049999.txt",
"invoice_0099999.txt",
"invoice_0149999.txt",
"invoice_0199999.txt",
"invoice_0249999.txt",
"invoice_0299999.txt",
"invoice_0349999.txt",
"invoice_0399999.txt",
"invoice_0449999.txt",
"invoice_0499999.txt",
"invoice_0549998.txt",
"invoice_0599998.txt",
"invoice_0649998.txt",
"invoice_0699998.txt",
"invoice_0749998.txt",
"invoice_0799998.txt",
"invoice_0849998.txt",
"invoice_0899998.txt",
"invoice_0949998.txt"
],
"samples_ms": [
85.3797,
14.3974,
13.200800000000001,
13.4454,
12.4215,
11.9434,
12.0694,
11.862400000000001,
12.014999999999999,
13.409600000000001,
12.745899999999999,
12.9113,
13.3632,
12.5779,
12.828199999999999,
13.0159,
12.9193,
12.663499999999999,
12.2875,
12.1379
],
"p50_ms": 12.745899999999999,
"p95_ms": 14.3974,
"first_sample_ms": 85.3797,
"first_sample_target_met": false,
"maximum_ms": 85.3797,
"samples_exceeding_target": 1,
"target_p95_ms": 50.0,
"target_met": true
},
"reopened_reader_first_names": {
"description": "each never-repeated filename on a newly opened reader, persisted statistics, no prewarm; includes open+search",
"includes_reader_open": true,
"queries": [
"invoice_0000001.txt",
"invoice_0050000.txt",
"invoice_0100000.txt",
"invoice_0150000.txt",
"invoice_0200000.txt",
"invoice_0250000.txt",
"invoice_0300000.txt",
"invoice_0350000.txt",
"invoice_0400000.txt",
"invoice_0450000.txt",
"invoice_0500000.txt",
"invoice_0549999.txt",
"invoice_0599999.txt",
"invoice_0649999.txt",
"invoice_0699999.txt",
"invoice_0749999.txt",
"invoice_0799999.txt",
"invoice_0849999.txt",
"invoice_0899999.txt",
"invoice_0949999.txt"
],
"samples_ms": [
100.0679,
19.9179,
17.9353,
16.6326,
16.3062,
16.3483,
16.519099999999998,
16.380499999999998,
19.8438,
16.5421,
15.4087,
31.4163,
14.6464,
14.1976,
14.2664,
15.4924,
16.0576,
14.2649,
13.9564,
16.1585
],
"p50_ms": 16.3062,
"p95_ms": 31.4163,
"first_sample_ms": 100.0679,
"first_sample_target_met": false,
"maximum_ms": 100.0679,
"samples_exceeding_target": 1,
"target_p95_ms": 50.0,
"target_met": true
},
"fixture": "synthetic SQLite metadata and text; no filesystem extraction or semantic inference",
"files": 1000000,
"documents": 100000,
"seed_seconds": 0.0,
"migration_seconds": 0.0118237,
"fixture_reused": true,
"database_bytes": 634449920,
"database_path": "target/benchmarks/full.db",
"architecture": "x86_64",
"os": "windows",
"cpu_only": true,
"queries": [
{
"query": "invoice_0999999.txt",
"samples": 20,
"p50_ms": 11.9678,
"first_query_ms": 13.4818,
"first_query_target_met": true,
"p95_ms": 13.218,
"target_p95_ms": 50.0,
"target_met": true
},
{
"query": "invoice_",
"samples": 20,
"p50_ms": 0.3231,
"first_query_ms": 3.2028999999999996,
"first_query_target_met": true,
"p95_ms": 0.4335,
"target_p95_ms": 50.0,
"target_met": true
},
{
"query": "marker0099999",
"samples": 20,
"p50_ms": 0.2697,
"first_query_ms": 2.5677,
"first_query_target_met": true,
"p95_ms": 0.5285,
"target_p95_ms": 150.0,
"target_met": true
},
{
"query": "cancellation",
"samples": 20,
"p50_ms": 95.59660000000001,
"first_query_ms": 99.10470000000001,
"first_query_target_met": true,
"p95_ms": 106.197,
"target_p95_ms": 150.0,
"target_met": true
}
]
}
crates/searchi-core/benchmarks/results/windows-cpu-full-suffix-anchor.json
+177
View File
@@ -0,0 +1,177 @@
{
"background_prepare_seconds": 4.8178274,
"reopened_background_prepare_seconds": 0.4527295,
"background_representative_queries": [
"invoice_0000000.txt",
"invoice_0999999.txt",
"Supplier",
"termination",
"cancellation"
],
"prepared_frequency_count": 1268,
"first_foreground_distinct_names": {
"description": "distinct foreground filenames; no benchmark query warmups; production background representative queries may overlap",
"includes_reader_open": false,
"queries": [
"invoice_0000000.txt",
"invoice_0049999.txt",
"invoice_0099999.txt",
"invoice_0149999.txt",
"invoice_0199999.txt",
"invoice_0249999.txt",
"invoice_0299999.txt",
"invoice_0349999.txt",
"invoice_0399999.txt",
"invoice_0449999.txt",
"invoice_0499999.txt",
"invoice_0549998.txt",
"invoice_0599998.txt",
"invoice_0649998.txt",
"invoice_0699998.txt",
"invoice_0749998.txt",
"invoice_0799998.txt",
"invoice_0849998.txt",
"invoice_0899998.txt",
"invoice_0949998.txt"
],
"samples_ms": [
136.6046,
43.111999999999995,
29.0124,
77.8574,
35.7493,
67.303,
92.4451,
91.5913,
29.5216,
66.8154,
27.6054,
366.9255,
28.3015,
70.88719999999999,
27.576900000000002,
69.3327,
27.537200000000002,
69.94200000000001,
26.957600000000003,
72.7744
],
"p50_ms": 66.8154,
"p95_ms": 136.6046,
"first_sample_ms": 136.6046,
"first_sample_target_met": false,
"maximum_ms": 366.9255,
"samples_exceeding_target": 11,
"target_p95_ms": 50.0,
"target_met": false
},
"reopened_reader_first_names": {
"description": "each never-repeated filename on a newly opened reader, persisted statistics, no prewarm; includes open+search",
"includes_reader_open": true,
"queries": [
"invoice_0000001.txt",
"invoice_0050000.txt",
"invoice_0100000.txt",
"invoice_0150000.txt",
"invoice_0200000.txt",
"invoice_0250000.txt",
"invoice_0300000.txt",
"invoice_0350000.txt",
"invoice_0400000.txt",
"invoice_0450000.txt",
"invoice_0500000.txt",
"invoice_0549999.txt",
"invoice_0599999.txt",
"invoice_0649999.txt",
"invoice_0699999.txt",
"invoice_0749999.txt",
"invoice_0799999.txt",
"invoice_0849999.txt",
"invoice_0899999.txt",
"invoice_0949999.txt"
],
"samples_ms": [
554.3185000000001,
129.17229999999998,
102.8,
102.7292,
101.44300000000001,
98.4861,
112.4415,
95.5881,
94.8502,
97.8484,
28.9007,
31.1041,
29.8992,
31.3256,
29.8718,
28.2319,
26.3245,
29.8964,
29.4154,
65.40020000000001
],
"p50_ms": 65.40020000000001,
"p95_ms": 129.17229999999998,
"first_sample_ms": 554.3185000000001,
"first_sample_target_met": false,
"maximum_ms": 554.3185000000001,
"samples_exceeding_target": 11,
"target_p95_ms": 50.0,
"target_met": false
},
"fixture": "synthetic SQLite metadata and text; no filesystem extraction or semantic inference",
"files": 1000000,
"documents": 100000,
"seed_seconds": 0.0,
"migration_seconds": 0.0049354,
"fixture_reused": true,
"database_bytes": 634449920,
"database_path": "target/benchmarks/full.db",
"architecture": "x86_64",
"os": "windows",
"cpu_only": true,
"queries": [
{
"query": "invoice_0999999.txt",
"samples": 20,
"p50_ms": 26.0376,
"first_query_ms": 28.9998,
"first_query_target_met": true,
"p95_ms": 30.639599999999998,
"target_p95_ms": 50.0,
"target_met": true
},
{
"query": "invoice_",
"samples": 20,
"p50_ms": 0.3764,
"first_query_ms": 73.6061,
"first_query_target_met": false,
"p95_ms": 0.8251000000000001,
"target_p95_ms": 50.0,
"target_met": true
},
{
"query": "marker0099999",
"samples": 20,
"p50_ms": 0.3811,
"first_query_ms": 9.8315,
"first_query_target_met": true,
"p95_ms": 0.5450999999999999,
"target_p95_ms": 150.0,
"target_met": true
},
{
"query": "cancellation",
"samples": 20,
"p50_ms": 97.4611,
"first_query_ms": 98.24900000000001,
"first_query_target_met": true,
"p95_ms": 112.3313,
"target_p95_ms": 150.0,
"target_met": true
}
]
}
crates/searchi-core/examples/performance.rs
+135 -4
View File
@@ -12,6 +12,7 @@ struct Measurements {
samples: usize,
p50_ms: f64,
first_query_ms: f64,
first_query_target_met: bool,
p95_ms: f64,
target_p95_ms: f64,
target_met: bool,
@@ -19,6 +20,12 @@ struct Measurements {
#[derive(Serialize)]
struct Report {
background_prepare_seconds: f64,
reopened_background_prepare_seconds: f64,
background_representative_queries: Vec<String>,
prepared_frequency_count: usize,
first_foreground_distinct_names: DistinctMeasurements,
reopened_reader_first_names: DistinctMeasurements,
fixture: &'static str,
files: usize,
documents: usize,
@@ -33,6 +40,84 @@ struct Report {
queries: Vec<Measurements>,
}
#[derive(Serialize)]
struct DistinctMeasurements {
description: &'static str,
includes_reader_open: bool,
queries: Vec<String>,
samples_ms: Vec<f64>,
p50_ms: f64,
p95_ms: f64,
first_sample_ms: f64,
first_sample_target_met: bool,
maximum_ms: f64,
samples_exceeding_target: usize,
target_p95_ms: f64,
target_met: bool,
}
fn distinct_names(
index: &SearchIndex,
database: &std::path::Path,
files: usize,
reopen: bool,
) -> Result<DistinctMeasurements> {
let count = 20.min(files / 2);
let mut queries = Vec::with_capacity(count);
let mut samples_ms = Vec::with_capacity(count);
for sample in 0..count {
// Two disjoint deterministic sets spread across the entire corpus.
// Neither set repeats the four-case exact filename at files-1.
let number = sample * (files - 2) / count + usize::from(reopen);
let query = format!("invoice_{number:07}.txt");
let start = Instant::now();
let results = if reopen {
SearchIndex::open(database)?.search(&query, 20)?
} else {
index.search(&query, 20)?
};
let elapsed = start.elapsed().as_secs_f64() * 1000.0;
let first = results
.first()
.context("distinct exact query returned no results")?;
if first.filename != query || first.match_kind != "exact_filename" {
bail!("distinct query lost exact filename priority");
}
queries.push(query);
samples_ms.push(elapsed);
}
let mut sorted = samples_ms.clone();
sorted.sort_by(f64::total_cmp);
let percentile = |fraction: f64| {
sorted[((sorted.len() as f64 * fraction).ceil() as usize).saturating_sub(1)]
};
let p95_ms = percentile(0.95);
eprintln!(
"{} distinct first filename queries (reopened={reopen}): p95 {p95_ms:.3} ms",
count
);
let first_sample_ms = samples_ms[0];
let samples_exceeding_target = samples_ms.iter().filter(|sample| **sample >= 50.0).count();
Ok(DistinctMeasurements {
description: if reopen {
"each never-repeated filename on a newly opened reader, persisted statistics, no prewarm; includes open+search"
} else {
"distinct foreground filenames; no benchmark query warmups; production background representative queries may overlap"
},
includes_reader_open: reopen,
queries,
samples_ms,
p50_ms: percentile(0.5),
p95_ms,
first_sample_ms,
first_sample_target_met: first_sample_ms < 50.0,
maximum_ms: sorted[sorted.len() - 1],
samples_exceeding_target,
target_p95_ms: 50.0,
target_met: p95_ms < 50.0,
})
}
fn main() -> Result<()> {
let arguments = env::args().skip(1).collect::<Vec<_>>();
let mut files = 10_000usize;
@@ -53,13 +138,12 @@ fn main() -> Result<()> {
option => bail!("unknown argument {option}"),
}
}
if files == 0
|| files > 2_000_000
if !(2..=2_000_000).contains(&files)
|| documents == 0
|| documents > files
|| !(5..=1000).contains(&samples)
{
bail!("files must be1..2000000, documents1..files, samples5..1000");
bail!("files must be2..2000000, documents1..files, samples5..1000");
}
if database.exists() && !reuse {
bail!(
@@ -135,7 +219,41 @@ fn main() -> Result<()> {
start.elapsed().as_secs_f64()
};
let database_bytes = fs::metadata(&database)?.len();
let index = SearchIndex::open(&database)?;
let mut index = SearchIndex::open(&database)?;
let connection = Connection::open(&database)?;
let mut statement=connection.prepare("SELECT filename FROM files WHERE id=(SELECT min(id) FROM files) OR id=(SELECT max(id) FROM files)")?;
let mut background_representative_queries = statement
.query_map([], |row| row.get::<_, String>(0))?
.collect::<rusqlite::Result<Vec<_>>>()?;
let words = background_representative_queries
.iter()
.flat_map(|name| {
name.split(|character: char| !character.is_alphabetic())
.filter(|word| word.chars().count() >= 3)
.take(3)
})
.map(str::to_owned)
.collect::<std::collections::BTreeSet<_>>();
background_representative_queries.extend(words);
let text: String = connection.query_row(
"SELECT substr(content,1,256) FROM files WHERE content!='' ORDER BY id LIMIT 1",
[],
|row| row.get(0),
)?;
background_representative_queries.extend(
text.split_whitespace()
.filter(|term| term.chars().all(char::is_alphanumeric))
.take(3)
.map(str::to_owned),
);
drop(statement);
drop(connection);
let preparation = Instant::now();
let prepared_frequency_count = index.prewarm()?;
let background_prepare_seconds = preparation.elapsed().as_secs_f64();
// Measure user-like first inputs before the old repeated-query cases.
let first_foreground_distinct_names = distinct_names(&index, &database, files, false)?;
let reopened_reader_first_names = distinct_names(&index, &database, files, true)?;
let exact = format!("invoice_{:07}.txt", files - 1);
let unique_content = format!("marker{:07}", documents - 1);
let cases = [
@@ -175,12 +293,25 @@ fn main() -> Result<()> {
samples,
p50_ms: percentile(0.5),
first_query_ms,
first_query_target_met: first_query_ms < target,
p95_ms,
target_p95_ms: target,
target_met: p95_ms < target,
});
}
// Measure restart preparation after foreground samples, so a second
// prewarm cannot secretly improve the first-input measurements above.
let mut restarted = SearchIndex::open(&database)?;
let restart_start = Instant::now();
restarted.prewarm()?;
let reopened_background_prepare_seconds = restart_start.elapsed().as_secs_f64();
let report = Report {
background_prepare_seconds,
reopened_background_prepare_seconds,
background_representative_queries,
prepared_frequency_count,
first_foreground_distinct_names,
reopened_reader_first_names,
fixture:
"synthetic SQLite metadata and text; no filesystem extraction or semantic inference",
files,
crates/searchi-core/src/lib.rs
+168 -16
View File
@@ -16,7 +16,7 @@ use std::{
};
use walkdir::WalkDir;
const SCHEMA_VERSION: i64 = 7;
const SCHEMA_VERSION: i64 = 9;
const MAX_RESULTS: usize = 1000;
#[derive(Debug, Default, Serialize)]
@@ -214,7 +214,18 @@ impl SearchIndex {
}
initialization.execute_batch(
"CREATE VIRTUAL TABLE IF NOT EXISTS names_vocab USING fts5vocab(names,'row');
PRAGMA user_version=7;",
CREATE TABLE IF NOT EXISTS name_frequencies(term TEXT PRIMARY KEY, doc INTEGER NOT NULL);
CREATE INDEX IF NOT EXISTS name_frequencies_popularity ON name_frequencies(doc DESC);
CREATE TABLE IF NOT EXISTS name_statistics_state(singleton INTEGER PRIMARY KEY CHECK(singleton=1), dirty INTEGER NOT NULL);
INSERT OR IGNORE INTO name_statistics_state VALUES(1,1);
CREATE TRIGGER IF NOT EXISTS files_statistics_insert AFTER INSERT ON files BEGIN
UPDATE name_statistics_state SET dirty=1 WHERE singleton=1; END;
CREATE TRIGGER IF NOT EXISTS files_statistics_delete AFTER DELETE ON files BEGIN
UPDATE name_statistics_state SET dirty=1 WHERE singleton=1; END;
CREATE TRIGGER IF NOT EXISTS files_statistics_update AFTER UPDATE OF filename,path ON files
WHEN new.filename!=old.filename OR new.path!=old.path BEGIN
UPDATE name_statistics_state SET dirty=1 WHERE singleton=1; END;
PRAGMA user_version=9;",
)?;
initialization.commit()?;
}
@@ -225,6 +236,89 @@ impl SearchIndex {
})
}
/// Refresh candidate-frequency statistics off the foreground search path.
/// Stale/missing estimates can affect speed, never literal-match correctness.
/// The app invokes this during background startup and after reconciliation.
pub fn prewarm(&mut self) -> Result<usize> {
let dirty: bool = self.connection.query_row(
"SELECT dirty FROM name_statistics_state WHERE singleton=1",
[],
|row| row.get(0),
)?;
if dirty {
let transaction = self
.connection
.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
transaction.execute("DELETE FROM name_frequencies", [])?;
transaction.execute(
"INSERT INTO name_frequencies(term,doc) SELECT term,doc FROM names_vocab",
[],
)?;
transaction.execute(
"UPDATE name_statistics_state SET dirty=0 WHERE singleton=1",
[],
)?;
transaction.commit()?;
}
let mut statement = self
.connection
.prepare("SELECT term,doc FROM name_frequencies ORDER BY doc DESC LIMIT 4096")?;
let rows = statement.query_map([], |row| {
Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?))
})?;
let mut frequencies = self.gram_frequencies.borrow_mut();
frequencies.clear();
for row in rows {
let (term, frequency) = row?;
frequencies.insert(term, frequency);
}
let cached = frequencies.len();
drop(frequencies);
// Exercise real query paths using representative corpus data, rather
// than paying postings, snippet and SQLite page initialization on input.
let mut representative = self.connection.prepare(
"SELECT filename FROM files WHERE id=(SELECT min(id) FROM files)
OR id=(SELECT max(id) FROM files)",
)?;
let names = representative
.query_map([], |row| row.get::<_, String>(0))?
.collect::<rusqlite::Result<Vec<_>>>()?;
for name in &names {
self.search(name, 20)?;
}
// Full names often anchor at sparse identifiers. Also exercise actual
// alphabetic name tokens to initialize common prefix/type postings.
let words = names
.iter()
.flat_map(|name| {
name.split(|character: char| !character.is_alphabetic())
.filter(|word| word.chars().count() >= 3)
.take(3)
})
.collect::<BTreeSet<_>>();
for word in words {
self.search(word, 20)?;
}
let text: Option<String> = self
.connection
.query_row(
"SELECT substr(content,1,256) FROM files WHERE content!='' ORDER BY id LIMIT 1",
[],
|row| row.get(0),
)
.optional()?;
if let Some(text) = text {
for term in text
.split_whitespace()
.filter(|term| term.chars().all(char::is_alphanumeric))
.take(3)
{
self.search(term, 20)?;
}
}
Ok(cached)
}
/// Reconcile one explicitly selected root. Symlinks are never followed.
/// Every file receives a filename record, including unsupported formats.
pub fn index_folder(&mut self, root: impl AsRef<Path>) -> Result<IndexReport> {
@@ -441,16 +535,17 @@ impl SearchIndex {
.modified_before_unix_ms
.map_or(Value::Null, Value::Integer),
];
let has_filters = filter_values.iter().any(|value| *value != Value::Null);
let mut results = Vec::new();
let mut selected = Vec::<i64>::new();
// Each lane filters its entire candidate set before LIMIT. There is no
// hidden pre-filter candidate cap that can hide later matching files.
let trigram = query.chars().count() >= 3;
// A long FTS phrase may walk common leading trigram postings even when
// later trigrams are rare. Use the rarest known trigram as a candidate
// superset, then verify the whole literal substring. The 256-term lookup
// budget can only enlarge candidates; it never caps results. Frequency
// lookups are cached because FTS vocabulary counts have a first-use cost.
// later trigrams are rare. Build a frequency-bounded positional anchor,
// then verify the whole literal substring. The 256-term lookup budget
// can only enlarge candidates; it never caps results. Frequencies come
// from persisted background statistics and a bounded reader cache.
let characters = query.chars().collect::<Vec<_>>();
let grams = characters
.windows(3)
@@ -461,7 +556,7 @@ impl SearchIndex {
.collect::<Vec<_>>();
let mut vocabulary = self
.connection
.prepare("SELECT doc FROM names_vocab WHERE term=?1")?;
.prepare("SELECT doc FROM name_frequencies WHERE term=?1")?;
let mut cache = self.gram_frequencies.borrow_mut();
if cache.len() + grams.len() > 4096 {
cache.clear();
@@ -485,12 +580,8 @@ impl SearchIndex {
.collect::<Result<Vec<_>>>()?;
rare_first.sort();
drop(cache);
let grams = rare_first
.into_iter()
.take(1)
.map(|(_, gram)| format!("\"{}\"", gram.replace('"', "\"\"")))
.collect::<Vec<_>>()
.join(" AND ");
let anchor = substring_anchor(&characters, &rare_first);
let grams = format!("\"{}\"", anchor.replace('"', "\"\""));
let mut lanes = vec![
(
"exact_filename",
@@ -585,10 +676,25 @@ impl SearchIndex {
} else {
"''"
};
let sql = format!(
"SELECT f.id,f.path,f.filename,{snippet},{rank},f.extraction_error FROM {from}
let sql = if kind == "content" && !has_filters {
// Rank compact row IDs before materializing paths and snippets.
// Filtered requests retain the complete filter-before-LIMIT lane
// below; this fast path is only for genuinely unfiltered search.
let excluded = exclusions.replace("f.id", "rowid");
format!("WITH hits AS MATERIALIZED (
SELECT rowid AS id,rank AS relevance FROM content_search
WHERE content_search MATCH ?1 {excluded}
ORDER BY rank,rowid LIMIT ?9)
SELECT f.id,f.path,f.filename,snippet(content_search,0,'','',' … ',24),hits.relevance,f.extraction_error
FROM hits CROSS JOIN files f ON f.id=hits.id
CROSS JOIN content_search ON content_search.rowid=hits.id
WHERE content_search MATCH ?1 ORDER BY hits.relevance,f.id")
} else {
format!(
"SELECT f.id,f.path,f.filename,{snippet},{rank},f.extraction_error FROM {from}
WHERE {condition} {filters} {exclusions} ORDER BY {ordering} LIMIT ?9"
);
)
};
let mut values = vec![Value::Text(expression), Value::Text(query.into())];
values.extend(filter_values.iter().cloned());
values.push(Value::Integer((limit - results.len()) as i64));
@@ -702,6 +808,52 @@ impl SearchIndex {
}
}
/// Build a bounded positional anchor around the rarest trigram. Extending a
/// phrase through common suffixes can force FTS to decode dense postings. Only
/// adjacent grams within a small frequency factor are admitted here. Every
/// literal match contains this anchor; SQL still verifies the entire substring.
fn substring_anchor(characters: &[char], frequencies: &[(i64, String)]) -> String {
let Some((minimum, gram)) = frequencies.first() else {
return String::new();
};
let center = characters
.windows(3)
.position(|window| window.iter().collect::<String>() == *gram)
.unwrap_or(0);
let mut start = center;
let mut end = center + 3;
if *minimum > 0 {
let maximum = minimum.saturating_mul(4).max(32);
let lookup = frequencies
.iter()
.map(|(frequency, gram)| (gram.as_str(), *frequency))
.collect::<HashMap<_, _>>();
let eligible = |index: usize| {
lookup
.get(
characters[index..index + 3]
.iter()
.collect::<String>()
.as_str(),
)
.is_some_and(|frequency| *frequency <= maximum)
};
loop {
let previous = (start, end);
if start > 0 && end - start < 32 && eligible(start - 1) {
start -= 1;
}
if end < characters.len() && end - start < 32 && eligible(end - 2) {
end += 1;
}
if previous == (start, end) {
break;
}
}
}
characters[start..end].iter().collect()
}
#[cfg(test)]
mod tests {
use super::*;
crates/searchi-core/tests/retrieval_lanes.rs
+103
View File
@@ -123,3 +123,106 @@ fn v4_migration_backfills_names_and_preserves_passage_ids() -> Result<()> {
assert_eq!(index.passage_batch(0, 1)?[0].id, 88);
Ok(())
}
#[test]
fn background_statistics_preserve_results_across_reopen_and_changes() -> Result<()> {
let root = tempfile::tempdir()?;
let db = root.path().join("index.db");
fs::write(root.path().join("invoice_1042.txt"), "billing evidence")?;
let mut index = SearchIndex::open(&db)?;
index.index_folder(root.path())?;
assert!(index.prewarm()? > 0);
drop(index);
let mut index = SearchIndex::open(&db)?;
assert_eq!(index.search("1042", 1)?[0].filename, "invoice_1042.txt");
fs::rename(
root.path().join("invoice_1042.txt"),
root.path().join("new_pattern.txt"),
)?;
index.index_folder(root.path())?;
// Stale statistics affect only candidate cost, never substring correctness.
assert_eq!(index.search("pattern", 1)?[0].filename, "new_pattern.txt");
assert!(index.search("1042", 1)?.is_empty());
index.prewarm()?;
assert_eq!(index.search("pattern", 1)?[0].filename, "new_pattern.txt");
Ok(())
}
#[test]
fn positional_candidates_preserve_repeated_identifiers_unicode_and_path_matches() -> Result<()> {
let root = tempfile::tempdir()?;
let nested = root.path().join("folder_0000001_rare");
fs::create_dir(&nested)?;
for (name, text) in [
("ζαζα_0000001_rare.txt", "first source"),
("ζαζα_0000101_rare.txt", "second source"),
("backup_ζαζα_0000001_rare.txt", "backup source"),
] {
fs::write(root.path().join(name), text)?;
}
fs::write(nested.join("other.md"), "path source")?;
let mut index = SearchIndex::open(":memory:")?;
index.index_folder(root.path())?;
index.prewarm()?;
let matches = index.search("0000001_rare", 20)?;
assert_eq!(matches.len(), 3);
assert_eq!(matches[0].match_kind, "filename");
assert_eq!(matches[1].match_kind, "filename");
assert_eq!(matches[2].match_kind, "path");
assert_eq!(
index.search("ζαζα_0000101_rare", 1)?[0].filename,
"ζαζα_0000101_rare.txt"
);
let filtered = SearchOptions {
folder: Some(nested),
extension: Some("md".into()),
..Default::default()
};
assert_eq!(
index.search_with_options("0000001_rare", 1, &filtered)?[0].filename,
"other.md"
);
assert!(index.search("0000010_rare", 20)?.is_empty());
Ok(())
}
#[test]
fn candidate_first_content_matches_complete_filtered_ranking_and_sources() -> Result<()> {
let root = tempfile::tempdir()?;
fs::write(root.path().join("needle.txt"), "needle once")?;
for number in 0..50 {
let text = format!(
"{} background evidence {}",
"needle ".repeat(number % 7 + 1),
"filler ".repeat(number % 5)
);
fs::write(root.path().join(format!("notes_{number:03}.md")), text)?;
}
let mut index = SearchIndex::open(":memory:")?;
index.index_folder(root.path())?;
index.prewarm()?;
let fast = index.search("needle", 20)?;
// This nonrestrictive filter selects the normal complete filtered lane.
let filtered = index.search_with_options(
"needle",
20,
&SearchOptions {
min_size: Some(0),
..Default::default()
},
)?;
assert_eq!(fast.len(), 20);
assert_eq!(filtered.len(), 20);
for (fast, filtered) in fast.iter().zip(&filtered) {
assert_eq!(fast.path, filtered.path);
assert_eq!(fast.match_kind, filtered.match_kind);
assert!((fast.score - filtered.score).abs() < 1e-12);
assert_eq!(fast.snippet, filtered.snippet);
assert_eq!(
fast.source.as_ref().unwrap().number,
filtered.source.as_ref().unwrap().number
);
}
assert_eq!(fast[0].match_kind, "filename");
Ok(())
}
crates/searchi-semantic/Cargo.toml
+9 -3
View File
@@ -3,13 +3,19 @@ name = "searchi-semantic"
version.workspace = true
edition.workspace = true
license.workspace = true
description = "Bounded offline CPU ONNX embedding adapter for Searchi"
description = "Bounded local Google EmbeddingGemma 2 worker adapter for Searchi"
[features]
default = []
minilm-proof = ["dep:ort", "dep:tokenizers"]
[dependencies]
anyhow.workspace = true
ort = { version = "=2.0.0-rc.13", default-features = false, features = ["std", "load-dynamic", "api-27"] }
ort = { version = "=2.0.0-rc.13", optional = true, default-features = false, features = ["std", "load-dynamic", "api-27"] }
sha2 = "0.10"
tokenizers = { version = "=0.22.2", default-features = false, features = ["fancy-regex"] }
tokenizers = { version = "=0.22.2", optional = true, default-features = false, features = ["fancy-regex"] }
serde.workspace = true
serde_json.workspace = true
[dev-dependencies]
serde_json.workspace = true
crates/searchi-semantic/native/CMakeLists.txt
+51
View File
@@ -0,0 +1,51 @@
cmake_minimum_required(VERSION 3.20)
project(searchi_gemma_worker LANGUAGES C CXX)
if(NOT DEFINED LLAMA_CPP_DIR OR NOT EXISTS "${LLAMA_CPP_DIR}/include/llama.h")
message(FATAL_ERROR "Supply -DLLAMA_CPP_DIR=<pinned patched llama.cpp b11461 source>")
endif()
set(CMAKE_CXX_STANDARD 17)
if(MSVC)
set(CMAKE_MSVC_RUNTIME_LIBRARY "MultiThreaded$<$<CONFIG:Debug>:Debug>" CACHE STRING "Standalone worker static MSVC runtime" FORCE)
endif()
set(BUILD_SHARED_LIBS OFF CACHE BOOL "Static local worker runtime" FORCE)
set(GGML_OPENMP OFF CACHE BOOL "Use bounded native llama thread pools without external OpenMP runtime" FORCE)
set(GGML_NATIVE OFF CACHE BOOL "Portable CPU kernels")
# GPU backends are explicit production build options. The Windows VM builds
# with all of these OFF; runtime device selection always retains CPU fallback.
set(GGML_CUDA OFF CACHE BOOL "Production CUDA backend")
set(GGML_VULKAN OFF CACHE BOOL "Production Vulkan backend")
set(GGML_METAL OFF CACHE BOOL "Production Metal backend")
set(SEARCHI_ACCELERATION "cpu" CACHE STRING "cpu, auto, metal, vulkan, or cuda")
if(SEARCHI_ACCELERATION STREQUAL "cpu")
set(GGML_CUDA OFF CACHE BOOL "CPU test build" FORCE)
set(GGML_VULKAN OFF CACHE BOOL "CPU test build" FORCE)
set(GGML_METAL OFF CACHE BOOL "CPU test build" FORCE)
elseif(SEARCHI_ACCELERATION STREQUAL "metal" OR (SEARCHI_ACCELERATION STREQUAL "auto" AND APPLE))
set(GGML_METAL ON CACHE BOOL "Production Metal backend" FORCE)
elseif(SEARCHI_ACCELERATION STREQUAL "cuda")
set(GGML_CUDA ON CACHE BOOL "Production CUDA backend" FORCE)
elseif(SEARCHI_ACCELERATION STREQUAL "vulkan")
set(GGML_VULKAN ON CACHE BOOL "Production Vulkan backend" FORCE)
elseif(SEARCHI_ACCELERATION STREQUAL "auto")
find_package(Vulkan QUIET)
if(Vulkan_FOUND)
set(GGML_VULKAN ON CACHE BOOL "Available production Vulkan backend" FORCE)
else()
message(STATUS "No production acceleration SDK detected; retaining CPU backend")
endif()
else()
message(FATAL_ERROR "Unknown SEARCHI_ACCELERATION=${SEARCHI_ACCELERATION}")
endif()
set(LLAMA_BUILD_COMMON OFF CACHE BOOL "No extra common runtime" FORCE)
set(LLAMA_BUILD_TOOLS OFF CACHE BOOL "No runtime tools" FORCE)
set(LLAMA_BUILD_SERVER OFF CACHE BOOL "No HTTP server" FORCE)
set(LLAMA_BUILD_EXAMPLES OFF CACHE BOOL "No examples" FORCE)
set(LLAMA_BUILD_TESTS OFF CACHE BOOL "No upstream tests in worker build" FORCE)
set(LLAMA_BUILD_MTMD OFF CACHE BOOL "Text worker; media encoder follows separately" FORCE)
set(LLAMA_OPENSSL OFF CACHE BOOL "No network runtime dependencies" FORCE)
add_subdirectory("${LLAMA_CPP_DIR}" llama-build)
add_executable(searchi-gemma-worker worker.cpp)
target_include_directories(searchi-gemma-worker PRIVATE "${LLAMA_CPP_DIR}/vendor")
target_link_libraries(searchi-gemma-worker PRIVATE llama)
target_compile_definitions(searchi-gemma-worker PRIVATE SEARCHI_LLAMA_TAG="b11461-no-logits-v1")
install(TARGETS searchi-gemma-worker RUNTIME DESTINATION bin)
crates/searchi-semantic/native/worker.cpp
+221
View File
@@ -0,0 +1,221 @@
// Local NDJSON helper around the pinned llama.cpp encoder. stdout is exclusively
// protocol output; llama diagnostics go to stderr. No model download or HTTP.
#include "llama.h"
#include "nlohmann/json.hpp"
#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <cstdio>
#include <iostream>
#include <memory>
#include <stdexcept>
#include <string>
#include <vector>
using json = nlohmann::json;
constexpr size_t max_request_bytes = 2 * 1024 * 1024;
constexpr size_t max_passage_bytes = 1024 * 1024;
constexpr size_t max_text_bytes = 65536;
constexpr size_t max_windows = 256;
constexpr int context_tokens = 2048;
constexpr int dimensions = 768;
static void require(bool condition, const char * message) {
if (!condition) throw std::runtime_error(message);
}
struct InferenceError : std::runtime_error { using std::runtime_error::runtime_error; };
static void require_inference(bool condition, const char * message) {
if (!condition) throw InferenceError(message);
}
class Encoder {
llama_model * model = nullptr;
llama_context * context = nullptr;
llama_batch_ext * batch = nullptr;
const llama_vocab * vocab = nullptr;
std::string model_path;
void release() {
if (batch) llama_batch_ext_free(batch);
if (context) llama_free(context);
if (model) llama_model_free(model);
batch = nullptr; context = nullptr; model = nullptr; vocab = nullptr;
}
bool open(bool gpu) {
auto mp = llama_model_default_params();
mp.n_gpu_layers = gpu ? 999 : 0;
mp.progress_callback = [](float,void *) {return true;};
model = llama_model_load_from_file(model_path.c_str(),mp);
if (!model) return false;
auto cp = llama_context_default_params();
cp.n_ctx = context_tokens; cp.n_batch = context_tokens; cp.n_ubatch = context_tokens;
cp.n_seq_max = 1; cp.kv_unified = true; cp.embeddings = true;
cp.pooling_type = LLAMA_POOLING_TYPE_MEAN;
cp.n_threads = threads; cp.n_threads_batch = threads;
context = llama_init_from_model(model,cp);
if (!context) {release(); return false;}
require(llama_model_n_embd_out(model) == dimensions,"EmbeddingGemma 2 output must have 768 dimensions");
vocab = llama_model_get_vocab(model);
batch = llama_batch_ext_init(context);
if (!batch) {release(); return false;}
device = gpu ? "gpu" : "cpu";
return true;
}
public:
std::string device;
int threads;
Encoder(const std::string & path, int cpu_threads, bool want_gpu) : model_path(path),threads(cpu_threads) {
llama_log_set([](ggml_log_level level, const char * text, void *) {
if (level >= GGML_LOG_LEVEL_WARN) std::fputs(text,stderr);
},nullptr);
llama_backend_init();
bool use_gpu = want_gpu && llama_supports_gpu_offload();
require(open(use_gpu) || (use_gpu && open(false)),"Cannot initialize local EmbeddingGemma 2 model/context");
}
~Encoder() {
release();
llama_backend_free();
}
bool healthy() const {return context != nullptr && batch != nullptr;}
void retry_cpu() {
require(device == "gpu","Only a GPU inference failure can request a CPU retry");
release();
require(open(false),"CPU fallback initialization failed after GPU inference failure");
}
std::vector<llama_token> tokenize(const std::string & text, bool specials) const {
std::vector<llama_token> tokens(text.size() + 8);
int count = llama_tokenize(vocab, text.data(), int(text.size()), tokens.data(), int(tokens.size()), specials, false);
if (count < 0) {
tokens.resize(size_t(-count));
count = llama_tokenize(vocab, text.data(), int(text.size()), tokens.data(), int(tokens.size()), specials, false);
}
require(count >= 0, "Tokenizer failed");
tokens.resize(size_t(count));
return tokens;
}
std::string decode(const std::vector<llama_token> & tokens, size_t count) const {
std::vector<char> text(max_passage_bytes + 64);
int length = llama_detokenize(vocab, tokens.data(), int(count), text.data(), int(text.size()), false, false);
require(length >= 0, "Detokenized passage exceeds byte budget");
return std::string(text.data(),size_t(length));
}
static std::string document_prefix(const std::string & title) {
require(title.size() <= 512, "Document title exceeds 512 UTF-8 bytes");
return "title: " + (title.empty() ? "none" : title) + " | text: ";
}
std::vector<float> embed(const std::string & text, int cap) {
require(!text.empty() && text.size() <= max_text_bytes, "Embedding text must be nonempty and at most 64 KiB");
auto tokens = tokenize(text,true);
require(!tokens.empty() && tokens.size() <= size_t(cap), "Prompt exceeds model token budget; split document windows first");
llama_batch_ext_clear(batch);
for (size_t i = 0; i < tokens.size(); ++i) {
int index = llama_batch_ext_add_token(batch,0,tokens[i]);
require(index >= 0, "Embedding batch is full");
llama_pos position[4] = {llama_pos(i),0,0,0};
require(llama_batch_ext_set_pos(batch,index,position), "Cannot set encoder position");
require(llama_batch_ext_set_output_embd(batch,index,true), "Cannot select embedding output");
}
require_inference(llama_process(context,LLAMA_PROCESS_TYPE_ENCODE,batch) == 0, "llama_process encoder failed");
const float * pooled = llama_get_embeddings_seq(context,0);
require_inference(pooled != nullptr, "Encoder returned no pooled vector");
double norm = 0;
for (int i = 0; i < dimensions; ++i) { require_inference(std::isfinite(pooled[i]), "Nonfinite encoder output"); norm += double(pooled[i])*pooled[i]; }
require_inference(norm > 0 && std::isfinite(norm), "Zero or nonfinite embedding norm");
std::vector<float> vector(dimensions);
for (int i = 0; i < dimensions; ++i) vector[i] = float(pooled[i]/std::sqrt(norm));
return vector;
}
json windows(const std::string & text, const std::string & title) const {
require(!text.empty() && text.size() <= max_passage_bytes, "Passage exceeds 1 MiB or is empty");
auto tokens = tokenize(text,false);
require(!tokens.empty(), "Passage has no tokens");
// Preserve actual source offsets, rather than inventing locations from
// a lossy decoded representation. Unexpected normalization is an error.
require(decode(tokens,tokens.size()) == text, "Tokenizer roundtrip changed source text; cannot verify source offsets");
size_t prompt = tokenize(document_prefix(title),true).size();
require(prompt + 132 < size_t(context_tokens), "Document prefix exceeds token budget");
size_t budget = size_t(context_tokens) - prompt - 4;
size_t stride = budget - 128;
json result = json::array();
for (size_t start = 0; start < tokens.size(); start += stride) {
require(result.size() < max_windows, "Passage exceeds 256 windows");
size_t end = std::min(start + budget,tokens.size());
size_t begin_byte = start ? decode(tokens,start).size() : 0;
size_t end_byte = decode(tokens,end).size();
// Token pieces can split a UTF-8 character. Enlarge to valid source
// boundaries; the four-token reserve covers boundary re-tokenization.
while (begin_byte > 0 && begin_byte < text.size() && (uint8_t(text[begin_byte]) & 0xc0) == 0x80) --begin_byte;
while (end_byte < text.size() && (uint8_t(text[end_byte]) & 0xc0) == 0x80) ++end_byte;
require(begin_byte < end_byte && end_byte <= text.size(), "Invalid source window byte range");
auto fragment = text.substr(begin_byte,end_byte-begin_byte);
require(fragment.size() + document_prefix(title).size() <= max_text_bytes, "Window exceeds text byte budget");
require(tokenize(document_prefix(title)+fragment,true).size() <= size_t(context_tokens), "Window re-tokenization exceeds context budget");
result.push_back({{"text",fragment},{"start_byte",begin_byte},{"end_byte",end_byte}});
if (end == tokens.size()) break;
}
return result;
}
json handle(const json & request) {
auto op = request.at("op").get<std::string>();
if (op == "windows") return {{"windows",windows(request.at("text").get<std::string>(),request.value("title",std::string()))}};
if (op == "query") return {{"vectors",json::array({embed("task: search result | query: "+request.at("text").get<std::string>(),512)})}};
if (op == "documents") {
const auto & texts = request.at("texts");
require(texts.is_array() && !texts.empty() && texts.size() <= 8,"Batch size must be between 1 and 8");
auto prefix = document_prefix(request.value("title",std::string()));
json vectors = json::array();
for (const auto & text : texts) vectors.push_back(embed(prefix+text.get<std::string>(),context_tokens));
return {{"vectors",vectors}};
}
throw std::runtime_error("Unknown embedding operation");
}
};
int main(int argc, char ** argv) {
try {
std::string path, device = "auto"; int threads = 1;
for (int i = 1; i < argc; ++i) {
require(i+1 < argc,"Worker options require values");
std::string option = argv[i++], value = argv[i];
if (option == "--model") path = value;
else if (option == "--threads") threads = std::stoi(value);
else if (option == "--device") device = value;
else throw std::runtime_error("Unknown worker option");
}
require(!path.empty() && threads >= 1 && threads <= 8,"Model path required; threads must be 1–8");
require(device == "cpu" || device == "auto","Device must be cpu or auto");
bool force_cpu = std::getenv("SEARCHI_FORCE_CPU") && std::string(std::getenv("SEARCHI_FORCE_CPU")) == "1";
Encoder encoder(path,threads,device == "auto" && !force_cpu);
std::cout << json({{"ready",true},{"model","google/embeddinggemma-2"},{"dimensions",dimensions},
{"device",encoder.device},{"threads",threads},{"runtime",SEARCHI_LLAMA_TAG}}).dump() << '\n' << std::flush;
for (;;) {
std::string line; char c;
while (std::cin.get(c) && c != '\n') {
require(line.size() < max_request_bytes,"Request exceeds 2 MiB"); line.push_back(c);
}
if (!std::cin && line.empty()) break;
json response;
try {
auto request = json::parse(line);
try {response = encoder.handle(request);}
catch (const InferenceError &) {
if (encoder.device != "gpu") throw;
encoder.retry_cpu();
response = encoder.handle(request); // Exactly one CPU retry.
response["fallback"] = "gpu_inference_failure";
}
response["ok"] = true;
}
catch (const std::exception & error) { response = {{"ok",false},{"error",error.what()},{"fatal",!encoder.healthy()}}; }
response["device"] = encoder.device;
auto encoded = response.dump();
require(encoded.size() <= 4*1024*1024,"Response exceeds 4 MiB");
std::cout << encoded << '\n' << std::flush;
if (!encoder.healthy()) return 1;
}
return 0;
} catch (const std::exception & error) {
std::cerr << "Searchi Gemma worker: " << error.what() << '\n';
return 1;
}
}
crates/searchi-semantic/src/gemma.rs
+398
View File
@@ -0,0 +1,398 @@
use crate::TextWindow;
use anyhow::{ensure, Context, Result};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use std::{
fs::File,
io::{BufRead, BufReader, Read, Write},
path::{Path, PathBuf},
process::{Child, ChildStdin, Command, Stdio},
sync::mpsc::{self, Receiver},
thread,
time::Duration,
};
pub const GEMMA_DIMENSIONS: usize = 768;
pub const GEMMA_FINGERPRINT: &str = "google/embeddinggemma-2:bfcd298762cc34d0357ece5ebdd31791a3a374d8:q8_0:llama-b11461-no-logits-v1:mean:l2:query512:doc2048:windows-overlap128:prompts-v1";
const MODEL_FILE: &str = "embeddinggemma-2-Q8_0.gguf";
const MODEL_BYTES: u64 = 309855456;
const MODEL_SHA256: &str = "2188ac1deca4b77dffefd603c2776a9d76d9d74ec01841392982ebb840b09135";
const MAX_REQUEST: usize = 2 * 1024 * 1024;
const MAX_RESPONSE: usize = 4 * 1024 * 1024;
#[derive(Debug, Clone, Copy)]
pub enum Device {
Auto,
Cpu,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct WorkerInfo {
pub ready: bool,
pub model: String,
pub dimensions: usize,
pub device: String,
pub threads: usize,
pub runtime: String,
}
/// Persistent isolated native worker. No HTTP or remote model fallback.
pub struct GemmaEmbedder {
child: Child,
stdin: ChildStdin,
responses: Receiver<std::result::Result<Value, String>>,
pub info: WorkerInfo,
healthy: bool,
force_cpu: bool,
model_dir: PathBuf,
worker_path: PathBuf,
}
impl Drop for GemmaEmbedder {
fn drop(&mut self) {
let _ = self.child.kill();
let _ = self.child.wait();
}
}
impl GemmaEmbedder {
/// Load the verified Google model and explicitly supplied local worker.
/// Production selects available acceleration with CPU fallback. Tests set
/// `SEARCHI_FORCE_CPU=1` or call `load_with_device(..., Device::Cpu)`.
pub fn load(
model_dir: impl AsRef<Path>,
worker: impl AsRef<Path>,
threads: usize,
) -> Result<Self> {
Self::load_with_device(model_dir, worker, threads, Device::Auto)
}
pub fn load_with_device(
model_dir: impl AsRef<Path>,
worker: impl AsRef<Path>,
threads: usize,
device: Device,
) -> Result<Self> {
ensure!(
(1..=8).contains(&threads),
"CPU threads must be between 1 and 8"
);
let model_dir = model_dir.as_ref().to_path_buf();
let model = model_dir.join(MODEL_FILE);
verify_model(&model)?;
let worker = worker
.as_ref()
.canonicalize()
.context("Cannot resolve local Gemma worker executable")?;
ensure!(
worker.is_file(),
"Gemma worker must be a local executable file"
);
let force_cpu = std::env::var("SEARCHI_FORCE_CPU").is_ok_and(|value| value == "1");
let cpu = force_cpu || matches!(device, Device::Cpu);
let mut command = Command::new(&worker);
command
.args(["--model"])
.arg(model)
.args([
"--threads",
&threads.to_string(),
"--device",
if cpu { "cpu" } else { "auto" },
])
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::inherit());
// Background helpers must not create a visible console on Windows.
#[cfg(windows)]
{
use std::os::windows::process::CommandExt;
command.creation_flags(0x08000000); // CREATE_NO_WINDOW
}
let mut child = command.spawn().context("Cannot start local Gemma worker")?;
let stdin = child.stdin.take().context("Worker stdin missing")?;
let stdout = child.stdout.take().context("Worker stdout missing")?;
let (sender, responses) = mpsc::sync_channel(1);
thread::spawn(move || {
let mut reader = BufReader::new(stdout);
loop {
let response = read_response(&mut reader);
let failed = response.is_err();
if sender.send(response).is_err() || failed {
break;
}
}
});
let ready = match responses.recv_timeout(Duration::from_secs(120)) {
Ok(Ok(value)) => value,
other => {
let _ = child.kill();
let _ = child.wait();
anyhow::bail!("Gemma worker did not initialize: {other:?}");
}
};
let info: WorkerInfo = match serde_json::from_value(ready) {
Ok(info) => info,
Err(error) => {
let _ = child.kill();
let _ = child.wait();
return Err(error).context("Invalid Gemma worker handshake");
}
};
if !info.ready
|| info.model != "google/embeddinggemma-2"
|| info.dimensions != GEMMA_DIMENSIONS
|| info.runtime != "b11461-no-logits-v1"
|| info.threads != threads
|| !["cpu", "gpu"].contains(&info.device.as_str())
|| (cpu && info.device != "cpu")
{
let _ = child.kill();
let _ = child.wait();
anyhow::bail!("Gemma worker handshake violates configured model/device contract");
}
Ok(Self {
child,
stdin,
responses,
info,
healthy: true,
force_cpu: cpu,
model_dir,
worker_path: worker,
})
}
fn request(&mut self, request: Value) -> Result<Value> {
ensure!(
self.healthy,
"Gemma worker session failed; reopen it explicitly"
);
let encoded = serde_json::to_vec(&request)?;
ensure!(
encoded.len() < MAX_REQUEST,
"Gemma request exceeds 2 MiB protocol budget"
);
let response = match self.exchange(&encoded) {
Ok(response) => response,
Err(error) if can_retry_cpu(&self.info.device, self.force_cpu) => {
// GPU runtimes can abort rather than return an encoder error.
// Restart once as CPU and retry this same bounded request.
let _ = self.child.wait();
let replacement = Self::load_with_device(
&self.model_dir,
&self.worker_path,
self.info.threads,
Device::Cpu,
)
.with_context(|| {
format!("CPU worker restart failed after GPU transport error: {error}")
})?;
*self = replacement;
self.exchange(&encoded)?
}
Err(error) => return Err(error),
};
if let Some(device) = response["device"].as_str() {
ensure!(
["cpu", "gpu"].contains(&device) && (!self.force_cpu || device == "cpu"),
"Worker response violates configured device policy"
);
self.info.device = device.into();
}
if response["fatal"] == true {
self.healthy = false;
self.info.ready = false;
let _ = self.child.kill();
}
ensure!(
response["ok"] == true,
"Gemma worker error: {}",
response["error"].as_str().unwrap_or("unknown error")
);
Ok(response)
}
fn exchange(&mut self, encoded: &[u8]) -> Result<Value> {
let sent = (|| {
self.stdin.write_all(encoded)?;
self.stdin.write_all(b"\n")?;
self.stdin.flush()
})()
.context("Cannot write local Gemma worker request");
if let Err(error) = sent {
self.healthy = false;
self.info.ready = false;
let _ = self.child.kill();
return Err(error);
}
let response = match self.responses.recv_timeout(Duration::from_secs(180)) {
Ok(Ok(value)) => value,
other => {
self.healthy = false;
self.info.ready = false;
let _ = self.child.kill();
anyhow::bail!("Gemma worker request failed: {other:?}");
}
};
Ok(response)
}
pub fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>> {
ensure!(
!text.trim().is_empty() && text.len() <= 1024 * 1024,
"Passage must be nonempty and at most 1 MiB"
);
let response = self.request(json!({"op":"windows","text":text,"title":title}))?;
let windows: Vec<TextWindow> = serde_json::from_value(response["windows"].clone())?;
ensure!(
!windows.is_empty() && windows.len() <= 256,
"Invalid Gemma window count"
);
for window in &windows {
ensure!(
text.get(window.start_byte..window.end_byte) == Some(window.text.as_str()),
"Gemma window does not match original source range"
);
}
ensure!(
windows
.last()
.is_some_and(|window| window.end_byte == text.len()),
"Gemma windows did not cover the passage tail"
);
Ok(windows)
}
pub fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result<Vec<Vec<f32>>> {
ensure!(
!texts.is_empty() && texts.len() <= 8,
"Document batch size must be between 1 and 8"
);
ensure!(
texts
.iter()
.all(|text| !text.trim().is_empty() && text.len() <= 65536),
"Document window must be nonempty and at most 64 KiB"
);
let response = self.request(json!({"op":"documents","texts":texts,"title":title}))?;
extract_vectors(response, texts.len())
}
pub fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
ensure!(
!query.trim().is_empty() && query.len() <= 4096,
"Query must be nonempty and at most 4096 UTF-8 bytes"
);
let response = self.request(json!({"op":"query","text":query}))?;
Ok(extract_vectors(response, 1)?.remove(0))
}
}
fn can_retry_cpu(actual_device: &str, force_cpu: bool) -> bool {
actual_device == "gpu" && !force_cpu
}
fn read_response(reader: &mut impl BufRead) -> std::result::Result<Value, String> {
let mut line = Vec::new();
match reader
.take((MAX_RESPONSE + 1) as u64)
.read_until(b'\n', &mut line)
{
Ok(0) => Err("Gemma worker closed its output".into()),
Ok(_) if line.len() > MAX_RESPONSE || line.last() != Some(&b'\n') => {
Err("Gemma worker response exceeds protocol byte budget".into())
}
Ok(_) => serde_json::from_slice(&line)
.map_err(|error| format!("Invalid Gemma worker JSON: {error}")),
Err(error) => Err(format!("Gemma worker read failed: {error}")),
}
}
fn extract_vectors(response: Value, count: usize) -> Result<Vec<Vec<f32>>> {
let vectors: Vec<Vec<f32>> = serde_json::from_value(response["vectors"].clone())?;
ensure!(
vectors.len() == count,
"Gemma returned incorrect embedding batch length"
);
for vector in &vectors {
ensure!(
vector.len() == GEMMA_DIMENSIONS && vector.iter().all(|value| value.is_finite()),
"Invalid Gemma embedding shape or values"
);
let norm = vector.iter().map(|v| v * v).sum::<f32>().sqrt();
ensure!(
(norm - 1.0).abs() < 1e-4,
"Gemma embedding is not L2 normalized"
);
}
Ok(vectors)
}
fn verify_model(path: &Path) -> Result<()> {
let file = File::open(path).with_context(|| {
format!(
"Missing Google EmbeddingGemma 2 model {}; run explicit model setup first",
path.display()
)
})?;
ensure!(
file.metadata()?.len() == MODEL_BYTES,
"EmbeddingGemma 2 model size mismatch"
);
let mut reader = BufReader::new(file);
let mut hash = Sha256::new();
let mut buffer = [0; 65536];
loop {
let count = reader.read(&mut buffer)?;
if count == 0 {
break;
}
hash.update(&buffer[..count]);
}
ensure!(
format!("{:x}", hash.finalize()) == MODEL_SHA256,
"EmbeddingGemma 2 model SHA256 mismatch"
);
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn missing_google_model_fails_without_starting_worker() {
assert!(GemmaEmbedder::load("missing-google-model", "missing-worker", 1).is_err());
}
#[test]
fn retry_policy_never_enables_gpu_or_retries_cpu_as_gpu() {
assert!(can_retry_cpu("gpu", false));
assert!(!can_retry_cpu("gpu", true));
assert!(!can_retry_cpu("cpu", false));
assert!(!can_retry_cpu("cpu", true));
}
#[test]
fn protocol_rejects_oversized_truncated_and_malformed_frames() {
assert!(read_response(&mut std::io::Cursor::new(vec![b'x'; MAX_RESPONSE + 1])).is_err());
assert!(read_response(&mut std::io::Cursor::new(b"{\"ok\":true}")).is_err());
assert!(read_response(&mut std::io::Cursor::new(b"not-json\n")).is_err());
assert_eq!(
read_response(&mut std::io::Cursor::new(b"{\"ok\":true}\n")).unwrap()["ok"],
true
);
}
#[test]
fn protocol_rejects_wrong_embedding_shape_count_and_normalization() {
assert!(extract_vectors(json!({"vectors":[]}), 1).is_err());
assert!(extract_vectors(json!({"vectors":[[1.0,0.0]]}), 1).is_err());
assert!(extract_vectors(json!({"vectors":[vec![0.0;GEMMA_DIMENSIONS]]}), 1).is_err());
let mut invalid = vec![0.0; GEMMA_DIMENSIONS];
invalid[0] = 1.0;
invalid[1] = 1.0;
assert!(extract_vectors(json!({"vectors":[invalid]}), 1).is_err());
}
}
crates/searchi-semantic/src/lib.rs
+17 -288
View File
@@ -1,290 +1,19 @@
//! Offline text embeddings using a pinned quantized multilingual model.
//!
//! Model provisioning is explicit and separate. This crate never downloads a
//! model or runtime and exposes no GPU execution-provider configuration.
//! Local Google EmbeddingGemma 2 inference through an isolated native worker.
//! The rejected MiniLM experiment is available only under `minilm-proof`.
mod gemma;
pub use gemma::{Device, GemmaEmbedder, WorkerInfo, GEMMA_DIMENSIONS, GEMMA_FINGERPRINT};
use std::{
fs::File,
io::{BufReader, Read},
path::{Path, PathBuf},
sync::Mutex,
/// A verified source byte range covered by a model token window.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct TextWindow {
pub text: String,
pub start_byte: usize,
pub end_byte: usize,
}
#[cfg(feature = "minilm-proof")]
mod minilm;
#[cfg(feature = "minilm-proof")]
pub use minilm::{
CpuEmbedder, DIMENSIONS, MAX_BATCH, MAX_TEXT_BYTES, MAX_TOKENS, MODEL_FINGERPRINT,
};
use anyhow::{bail, ensure, Context, Result};
use ort::{ep, session::Session, value::Tensor};
use sha2::{Digest, Sha256};
use tokenizers::{PaddingDirection, PaddingParams, Tokenizer, TruncationParams};
/// Embedding dimensionality; stored vectors must carry the model fingerprint.
pub const DIMENSIONS: usize = 384;
pub const MAX_TOKENS: usize = 128;
pub const MAX_BATCH: usize = 8;
pub const MAX_TEXT_BYTES: usize = 65536;
/// Includes model revision, quantization, pooling, normalization and token cap.
pub const MODEL_FINGERPRINT: &str =
"minilm-multilingual:2c4055b12046f11709e9df2c122e59ffbdc2f900:int8:mean:l2:128";
const ARTIFACTS: [(&str, u64, &str); 2] = [
(
"onnx/model_quantized.onnx",
118308126,
"66fc00f5f29afcaff34092e1bdd20008ca3918265a82fb9695a551e510cc4ebc",
),
(
"tokenizer.json",
17082913,
"b60b6b43406a48bf3638526314f3d232d97058bc93472ff2de930d43686fa441",
),
];
// ONNX Runtime's library and environment are process-global. Prevent concurrent
// initializers from selecting different libraries or silently overriding state.
static RUNTIME_PATH: Mutex<Option<PathBuf>> = Mutex::new(None);
/// Local, mutable inference session. Call from a bounded background worker.
pub struct CpuEmbedder {
session: Session,
tokenizer: Tokenizer,
}
impl CpuEmbedder {
/// Verify local artifacts, then load the explicitly supplied runtime.
///
/// The runtime must be initialized through this adapter before other `ort`
/// code in this process. Threads are limited to 1–8; default caller policy
/// should use one thread in a VM. Missing or damaged assets are errors.
pub fn load(
model_dir: impl AsRef<Path>,
runtime_library: impl AsRef<Path>,
threads: usize,
) -> Result<Self> {
ensure!(
(1..=8).contains(&threads),
"CPU threads must be between 1 and 8"
);
let model_dir = model_dir.as_ref();
verify_artifacts(model_dir)?;
let runtime_path = runtime_library
.as_ref()
.canonicalize()
.context("Cannot resolve explicitly supplied ONNX Runtime library")?;
ensure!(
runtime_path.is_file(),
"ONNX Runtime library must be a file"
);
{
let mut initialized = RUNTIME_PATH
.lock()
.map_err(|_| anyhow::anyhow!("Runtime initialization lock poisoned"))?;
if let Some(existing) = initialized.as_ref() {
ensure!(
*existing == runtime_path,
"ONNX Runtime already initialized with a different library"
);
} else {
let committed = ort::init_from(&runtime_path)?
.with_telemetry(false)
.with_execution_providers([ep::CPU::default().build()])
.commit();
ensure!(
committed,
"ONNX Runtime environment was already initialized outside the CPU adapter"
);
*initialized = Some(runtime_path);
}
}
let session = Session::builder()?
.with_execution_providers([ep::CPU::default().build()])
.map_err(|e| anyhow::anyhow!("CPU session configuration: {e}"))?
.with_intra_threads(threads)
.map_err(|e| anyhow::anyhow!("CPU thread configuration: {e}"))?
.with_inter_threads(1)
.map_err(|e| anyhow::anyhow!("Inter-thread configuration: {e}"))?
.with_parallel_execution(false)
.map_err(|e| anyhow::anyhow!("Sequential session configuration: {e}"))?
.commit_from_file(model_dir.join("onnx/model_quantized.onnx"))?;
let mut tokenizer = Tokenizer::from_file(model_dir.join("tokenizer.json"))
.map_err(|error| anyhow::anyhow!("Cannot load tokenizer: {error}"))?;
tokenizer
.with_truncation(Some(TruncationParams {
max_length: MAX_TOKENS,
..Default::default()
}))
.map_err(|error| anyhow::anyhow!("Cannot configure tokenizer: {error}"))?;
tokenizer.with_padding(Some(PaddingParams {
pad_id: 1,
pad_token: "<pad>".into(),
..Default::default()
}));
Ok(Self { session, tokenizer })
}
/// Produce finite, L2-normalized vectors from genuine local model inference.
///
/// Inputs above 128 tokens are truncated. Production indexing should chunk
/// documents first and keep chunk locations; never silently embed entire PDFs.
pub fn encode(&mut self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
validate_texts(texts)?;
// Encode sequentially so tokenizers cannot create an unbounded global
// Rayon worker pool independently of the caller's inference budget.
let mut encoded = texts
.iter()
.map(|text| self.tokenizer.encode(*text, true))
.collect::<std::result::Result<Vec<_>, _>>()
.map_err(|error| anyhow::anyhow!("Tokenization failed: {error}"))?;
let sequence_len = encoded.iter().map(|item| item.len()).max().unwrap_or(0);
for item in &mut encoded {
item.pad(sequence_len, 1, 0, "<pad>", PaddingDirection::Right);
}
ensure!(
sequence_len > 0 && sequence_len <= MAX_TOKENS,
"Invalid token sequence length"
);
let ids: Vec<i64> = encoded
.iter()
.flat_map(|item| item.get_ids().iter().map(|&v| i64::from(v)))
.collect();
let masks: Vec<i64> = encoded
.iter()
.flat_map(|item| item.get_attention_mask().iter().map(|&v| i64::from(v)))
.collect();
let types: Vec<i64> = encoded
.iter()
.flat_map(|item| item.get_type_ids().iter().map(|&v| i64::from(v)))
.collect();
let shape = [texts.len(), sequence_len];
let outputs = self.session.run(ort::inputs![
"input_ids" => Tensor::from_array((shape, ids))?,
"attention_mask" => Tensor::from_array((shape, masks.clone()))?,
"token_type_ids" => Tensor::from_array((shape, types))?
])?;
let (output_shape, hidden) = outputs[0].try_extract_tensor::<f32>()?;
ensure!(
**output_shape == [texts.len() as i64, sequence_len as i64, DIMENSIONS as i64],
"Unexpected embedding output shape"
);
pool_and_normalize(hidden, &masks, texts.len(), sequence_len)
}
}
fn validate_texts(texts: &[&str]) -> Result<()> {
ensure!(
!texts.is_empty() && texts.len() <= MAX_BATCH,
"Batch size must be between 1 and {MAX_BATCH}"
);
ensure!(
texts
.iter()
.all(|text| !text.trim().is_empty() && text.len() <= MAX_TEXT_BYTES),
"Text must be nonempty and at most {MAX_TEXT_BYTES} UTF-8 bytes"
);
Ok(())
}
fn verify_artifacts(model_dir: &Path) -> Result<()> {
for (relative, bytes, expected) in ARTIFACTS {
let path = model_dir.join(relative);
let file = File::open(&path).with_context(|| {
format!(
"Missing local model artifact {}; run the explicit model downloader first",
path.display()
)
})?;
ensure!(
file.metadata()?.len() == bytes,
"Model artifact size mismatch: {}",
path.display()
);
let mut reader = BufReader::new(file);
let mut hash = Sha256::new();
let mut buffer = [0_u8; 65536];
loop {
let read = reader.read(&mut buffer)?;
if read == 0 {
break;
}
hash.update(&buffer[..read]);
}
let actual = format!("{:x}", hash.finalize());
ensure!(
actual == expected,
"Model artifact SHA256 mismatch: {}",
path.display()
);
}
Ok(())
}
fn pool_and_normalize(
hidden: &[f32],
masks: &[i64],
batch: usize,
tokens: usize,
) -> Result<Vec<Vec<f32>>> {
ensure!(
hidden.len() == batch * tokens * DIMENSIONS && masks.len() == batch * tokens,
"Invalid pooling input dimensions"
);
let mut vectors = Vec::with_capacity(batch);
for item in 0..batch {
let mut vector = vec![0_f32; DIMENSIONS];
let mut count = 0;
for token in 0..tokens {
if masks[item * tokens + token] == 0 {
continue;
}
count += 1;
let offset = (item * tokens + token) * DIMENSIONS;
for dimension in 0..DIMENSIONS {
vector[dimension] += hidden[offset + dimension];
}
}
ensure!(count > 0, "Cannot pool a fully masked sequence");
for value in &mut vector {
*value /= count as f32;
}
let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
if !norm.is_finite() || norm <= f32::EPSILON {
bail!("Model produced nonfinite or zero embeddings");
}
for value in &mut vector {
*value /= norm;
}
ensure!(
vector.iter().all(|value| value.is_finite()),
"Model produced nonfinite embeddings"
);
vectors.push(vector);
}
Ok(vectors)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn padding_is_excluded_and_output_is_normalized() {
let mut hidden = vec![0.0; 2 * DIMENSIONS];
hidden[0] = 3.0;
hidden[1] = 4.0;
hidden[DIMENSIONS] = 999.0;
let pooled = pool_and_normalize(&hidden, &[1, 0], 1, 2).unwrap();
assert!((pooled[0][0] - 0.6).abs() < 1e-6);
assert!((pooled[0][1] - 0.8).abs() < 1e-6);
}
#[test]
fn rejects_unbounded_inputs_and_invalid_outputs() {
assert!(validate_texts(&[]).is_err());
assert!(validate_texts(&["x"; MAX_BATCH + 1]).is_err());
assert!(validate_texts(&[" "]).is_err());
assert!(validate_texts(&[&"x".repeat(MAX_TEXT_BYTES + 1)]).is_err());
assert!(pool_and_normalize(&vec![0.0; DIMENSIONS], &[0], 1, 1).is_err());
assert!(pool_and_normalize(&vec![f32::NAN; DIMENSIONS], &[1], 1, 1).is_err());
}
#[test]
fn missing_model_does_not_load_or_download_runtime() {
assert!(CpuEmbedder::load("this-model-does-not-exist", "missing-runtime", 1).is_err());
}
}
crates/searchi-semantic/src/minilm.rs
+347
View File
@@ -0,0 +1,347 @@
//! Offline text embeddings using a pinned quantized multilingual model.
//!
//! Model provisioning is explicit and separate. This crate never downloads a
//! model or runtime and exposes no GPU execution-provider configuration.
use std::{
fs::File,
io::{BufReader, Read},
path::{Path, PathBuf},
sync::Mutex,
};
use anyhow::{bail, ensure, Context, Result};
use ort::{ep, session::Session, value::Tensor};
use sha2::{Digest, Sha256};
use tokenizers::{PaddingDirection, PaddingParams, Tokenizer, TruncationParams};
/// Embedding dimensionality; stored vectors must carry the model fingerprint.
pub const DIMENSIONS: usize = 384;
pub const MAX_TOKENS: usize = 128;
pub const MAX_BATCH: usize = 8;
pub const MAX_TEXT_BYTES: usize = 65536;
pub const MAX_PASSAGE_BYTES: usize = 1024 * 1024;
pub const MAX_WINDOWS: usize = 256;
pub const WINDOW_OVERLAP_TOKENS: usize = 24;
/// A source byte range covered by one overlapping model token window.
use crate::TextWindow;
/// Includes model revision, quantization, pooling, normalization and token cap.
pub const MODEL_FINGERPRINT: &str =
"minilm-multilingual:2c4055b12046f11709e9df2c122e59ffbdc2f900:int8:mean:l2:128";
const ARTIFACTS: [(&str, u64, &str); 2] = [
(
"onnx/model_quantized.onnx",
118308126,
"66fc00f5f29afcaff34092e1bdd20008ca3918265a82fb9695a551e510cc4ebc",
),
(
"tokenizer.json",
17082913,
"b60b6b43406a48bf3638526314f3d232d97058bc93472ff2de930d43686fa441",
),
];
// ONNX Runtime's library and environment are process-global. Prevent concurrent
// initializers from selecting different libraries or silently overriding state.
static RUNTIME_PATH: Mutex<Option<PathBuf>> = Mutex::new(None);
/// Local, mutable inference session. Call from a bounded background worker.
pub struct CpuEmbedder {
session: Session,
tokenizer: Tokenizer,
}
impl CpuEmbedder {
/// Split a bounded passage into overlapping token windows, retaining its tail.
/// Special tokens have no source range; source offsets remain UTF-8 byte
/// offsets into the original passage for accurate preview highlighting.
pub fn split_windows(&self, text: &str) -> Result<Vec<TextWindow>> {
ensure!(
!text.trim().is_empty() && text.len() <= MAX_PASSAGE_BYTES,
"Passage must be nonempty and at most {MAX_PASSAGE_BYTES} UTF-8 bytes"
);
let mut tokenizer = self.tokenizer.clone();
tokenizer
.with_truncation(Some(TruncationParams {
max_length: MAX_TOKENS,
stride: WINDOW_OVERLAP_TOKENS,
..Default::default()
}))
.map_err(|error| anyhow::anyhow!("Window tokenizer configuration: {error}"))?;
tokenizer.with_padding(None);
let encoded = tokenizer
.encode(text, true)
.map_err(|error| anyhow::anyhow!("Window tokenization failed: {error}"))?;
ensure!(
encoded.get_overflowing().len() < MAX_WINDOWS,
"Passage exceeds {MAX_WINDOWS} model windows; split it before embedding"
);
let mut windows = Vec::new();
for item in std::iter::once(&encoded).chain(encoded.get_overflowing()) {
let offsets: Vec<_> = item
.get_offsets()
.iter()
.zip(item.get_special_tokens_mask())
.filter(|((start, end), special)| **special == 0 && end > start)
.map(|(offset, _)| *offset)
.collect();
let start_byte = offsets.iter().map(|(start, _)| *start).min().unwrap_or(0);
let end_byte = offsets.iter().map(|(_, end)| *end).max().unwrap_or(0);
ensure!(start_byte < end_byte, "Token window has no source text");
let chunk = text
.get(start_byte..end_byte)
.context("Invalid UTF-8 tokenizer offsets")?;
ensure!(
chunk.len() <= MAX_TEXT_BYTES,
"Token window exceeds embedding text byte budget"
);
windows.push(TextWindow {
text: chunk.to_owned(),
start_byte,
end_byte,
});
}
Ok(windows)
}
/// Verify local artifacts, then load the explicitly supplied runtime.
///
/// The runtime must be initialized through this adapter before other `ort`
/// code in this process. Threads are limited to 1–8; default caller policy
/// should use one thread in a VM. Missing or damaged assets are errors.
pub fn load(
model_dir: impl AsRef<Path>,
runtime_library: impl AsRef<Path>,
threads: usize,
) -> Result<Self> {
ensure!(
(1..=8).contains(&threads),
"CPU threads must be between 1 and 8"
);
let model_dir = model_dir.as_ref();
verify_artifacts(model_dir)?;
let runtime_path = runtime_library
.as_ref()
.canonicalize()
.context("Cannot resolve explicitly supplied ONNX Runtime library")?;
ensure!(
runtime_path.is_file(),
"ONNX Runtime library must be a file"
);
{
let mut initialized = RUNTIME_PATH
.lock()
.map_err(|_| anyhow::anyhow!("Runtime initialization lock poisoned"))?;
if let Some(existing) = initialized.as_ref() {
ensure!(
*existing == runtime_path,
"ONNX Runtime already initialized with a different library"
);
} else {
let committed = ort::init_from(&runtime_path)?
.with_telemetry(false)
.with_execution_providers([ep::CPU::default().build()])
.commit();
ensure!(
committed,
"ONNX Runtime environment was already initialized outside the CPU adapter"
);
*initialized = Some(runtime_path);
}
}
let session = Session::builder()?
.with_execution_providers([ep::CPU::default().build()])
.map_err(|e| anyhow::anyhow!("CPU session configuration: {e}"))?
.with_intra_threads(threads)
.map_err(|e| anyhow::anyhow!("CPU thread configuration: {e}"))?
.with_inter_threads(1)
.map_err(|e| anyhow::anyhow!("Inter-thread configuration: {e}"))?
.with_parallel_execution(false)
.map_err(|e| anyhow::anyhow!("Sequential session configuration: {e}"))?
.commit_from_file(model_dir.join("onnx/model_quantized.onnx"))?;
let mut tokenizer = Tokenizer::from_file(model_dir.join("tokenizer.json"))
.map_err(|error| anyhow::anyhow!("Cannot load tokenizer: {error}"))?;
tokenizer
.with_truncation(Some(TruncationParams {
max_length: MAX_TOKENS,
..Default::default()
}))
.map_err(|error| anyhow::anyhow!("Cannot configure tokenizer: {error}"))?;
tokenizer.with_padding(Some(PaddingParams {
pad_id: 1,
pad_token: "<pad>".into(),
..Default::default()
}));
Ok(Self { session, tokenizer })
}
/// Produce finite, L2-normalized vectors from genuine local model inference.
///
/// Inputs above 128 tokens are truncated. Production indexing should chunk
/// documents first and keep chunk locations; never silently embed entire PDFs.
pub fn encode(&mut self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
validate_texts(texts)?;
// Encode sequentially so tokenizers cannot create an unbounded global
// Rayon worker pool independently of the caller's inference budget.
let mut encoded = texts
.iter()
.map(|text| self.tokenizer.encode(*text, true))
.collect::<std::result::Result<Vec<_>, _>>()
.map_err(|error| anyhow::anyhow!("Tokenization failed: {error}"))?;
let sequence_len = encoded.iter().map(|item| item.len()).max().unwrap_or(0);
for item in &mut encoded {
item.pad(sequence_len, 1, 0, "<pad>", PaddingDirection::Right);
}
ensure!(
sequence_len > 0 && sequence_len <= MAX_TOKENS,
"Invalid token sequence length"
);
let ids: Vec<i64> = encoded
.iter()
.flat_map(|item| item.get_ids().iter().map(|&v| i64::from(v)))
.collect();
let masks: Vec<i64> = encoded
.iter()
.flat_map(|item| item.get_attention_mask().iter().map(|&v| i64::from(v)))
.collect();
let types: Vec<i64> = encoded
.iter()
.flat_map(|item| item.get_type_ids().iter().map(|&v| i64::from(v)))
.collect();
let shape = [texts.len(), sequence_len];
let outputs = self.session.run(ort::inputs![
"input_ids" => Tensor::from_array((shape, ids))?,
"attention_mask" => Tensor::from_array((shape, masks.clone()))?,
"token_type_ids" => Tensor::from_array((shape, types))?
])?;
let (output_shape, hidden) = outputs[0].try_extract_tensor::<f32>()?;
ensure!(
**output_shape == [texts.len() as i64, sequence_len as i64, DIMENSIONS as i64],
"Unexpected embedding output shape"
);
pool_and_normalize(hidden, &masks, texts.len(), sequence_len)
}
}
fn validate_texts(texts: &[&str]) -> Result<()> {
ensure!(
!texts.is_empty() && texts.len() <= MAX_BATCH,
"Batch size must be between 1 and {MAX_BATCH}"
);
ensure!(
texts
.iter()
.all(|text| !text.trim().is_empty() && text.len() <= MAX_TEXT_BYTES),
"Text must be nonempty and at most {MAX_TEXT_BYTES} UTF-8 bytes"
);
Ok(())
}
fn verify_artifacts(model_dir: &Path) -> Result<()> {
for (relative, bytes, expected) in ARTIFACTS {
let path = model_dir.join(relative);
let file = File::open(&path).with_context(|| {
format!(
"Missing local model artifact {}; run the explicit model downloader first",
path.display()
)
})?;
ensure!(
file.metadata()?.len() == bytes,
"Model artifact size mismatch: {}",
path.display()
);
let mut reader = BufReader::new(file);
let mut hash = Sha256::new();
let mut buffer = [0_u8; 65536];
loop {
let read = reader.read(&mut buffer)?;
if read == 0 {
break;
}
hash.update(&buffer[..read]);
}
let actual = format!("{:x}", hash.finalize());
ensure!(
actual == expected,
"Model artifact SHA256 mismatch: {}",
path.display()
);
}
Ok(())
}
fn pool_and_normalize(
hidden: &[f32],
masks: &[i64],
batch: usize,
tokens: usize,
) -> Result<Vec<Vec<f32>>> {
ensure!(
hidden.len() == batch * tokens * DIMENSIONS && masks.len() == batch * tokens,
"Invalid pooling input dimensions"
);
let mut vectors = Vec::with_capacity(batch);
for item in 0..batch {
let mut vector = vec![0_f32; DIMENSIONS];
let mut count = 0;
for token in 0..tokens {
if masks[item * tokens + token] == 0 {
continue;
}
count += 1;
let offset = (item * tokens + token) * DIMENSIONS;
for dimension in 0..DIMENSIONS {
vector[dimension] += hidden[offset + dimension];
}
}
ensure!(count > 0, "Cannot pool a fully masked sequence");
for value in &mut vector {
*value /= count as f32;
}
let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
if !norm.is_finite() || norm <= f32::EPSILON {
bail!("Model produced nonfinite or zero embeddings");
}
for value in &mut vector {
*value /= norm;
}
ensure!(
vector.iter().all(|value| value.is_finite()),
"Model produced nonfinite embeddings"
);
vectors.push(vector);
}
Ok(vectors)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn padding_is_excluded_and_output_is_normalized() {
let mut hidden = vec![0.0; 2 * DIMENSIONS];
hidden[0] = 3.0;
hidden[1] = 4.0;
hidden[DIMENSIONS] = 999.0;
let pooled = pool_and_normalize(&hidden, &[1, 0], 1, 2).unwrap();
assert!((pooled[0][0] - 0.6).abs() < 1e-6);
assert!((pooled[0][1] - 0.8).abs() < 1e-6);
}
#[test]
fn rejects_unbounded_inputs_and_invalid_outputs() {
assert!(validate_texts(&[]).is_err());
assert!(validate_texts(&["x"; MAX_BATCH + 1]).is_err());
assert!(validate_texts(&[" "]).is_err());
assert!(validate_texts(&[&"x".repeat(MAX_TEXT_BYTES + 1)]).is_err());
assert!(pool_and_normalize(&vec![0.0; DIMENSIONS], &[0], 1, 1).is_err());
assert!(pool_and_normalize(&vec![f32::NAN; DIMENSIONS], &[1], 1, 1).is_err());
}
#[test]
fn missing_model_does_not_load_or_download_runtime() {
assert!(CpuEmbedder::load("this-model-does-not-exist", "missing-runtime", 1).is_err());
}
}
crates/searchi-semantic/tests/cpu_model.rs
+1
View File
@@ -1,4 +1,5 @@
//! Opt-in real-artifact test. Default tests never acquire model/runtime assets.
#![cfg(feature = "minilm-proof")]
use searchi_semantic::{CpuEmbedder, DIMENSIONS};
use std::path::PathBuf;
crates/searchi-semantic/tests/fixtures/LICENSE-DigUp.txt
Vendored
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2026 Abdur Rahim
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
crates/searchi-semantic/tests/fixtures/digup-text-reference.json
Vendored
+3858
View File
File diff suppressed because it is too large. Load diff
crates/searchi-semantic/tests/gemma_cpu.rs
+102
View File
@@ -0,0 +1,102 @@
use anyhow::Result;
use searchi_semantic::{Device, GemmaEmbedder, GEMMA_DIMENSIONS};
#[test]
#[ignore = "Requires explicitly provisioned Google model and CPU worker"]
fn google_cpu_real_embeddings_and_verified_tail_windows() -> Result<()> {
let model = std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Set SEARCHI_GEMMA_MODEL_DIR");
let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Set SEARCHI_GEMMA_WORKER");
let mut embedder = GemmaEmbedder::load_with_device(model, worker, 1, Device::Cpu)?;
assert_eq!(embedder.info.device, "cpu");
assert_eq!(embedder.info.dimensions, GEMMA_DIMENSIONS);
let reference: serde_json::Value =
serde_json::from_str(include_str!("fixtures/digup-text-reference.json"))?;
let long_reference=(1..=60).map(|i|format!("Note {i}: the review moved to room {} because the projector in the main hall failed again, and the team agreed to send the slides the night before.",i%7+2)).collect::<Vec<_>>().join(" ");
let mut made = Vec::new();
for (name, query) in [
("query", "payment declined error"),
("query-bn", "ইলিশ মাছ রান্না"),
("query-ar", "متى تغلق المكتبة"),
] {
made.push((name, embedder.embed_query(query)?));
}
made.push(("document",embedder.embed_documents(&["Flight to Tokyo Haneda on November 3. Then the train to Kamakura to see the Great Buddha, and back for dinner in Shibuya."],"Tokyo itinerary")?.remove(0)));
made.push((
"document-long",
embedder.embed_documents(&[&long_reference], "")?.remove(0),
));
for (name, vector) in made {
let expected = reference["vectors"][name].as_array().unwrap();
let cosine: f64 = vector
.iter()
.zip(expected)
.map(|(a, b)| f64::from(*a) * b.as_f64().unwrap())
.sum();
println!("DigUp text reference {name}: cosine {cosine:.8}");
assert!(
cosine >= 0.999,
"Google CPU reference parity failed for {name}: {cosine}"
);
}
let documents = [
"The supplier agreement may be ended with sixty days written notice.",
"Business train tickets and hotel stays can be reimbursed by submitting receipts.",
"Die Gartenplanung sieht ein Hochbeet für Gemüse und eine Regentonne vor.",
];
let vectors = embedder.embed_documents(&documents, "none")?;
for (query, expected) in [
("Wie beenden wir die Zusammenarbeit mit dem Lieferanten?", 0),
("How do I claim my accommodation and railway expenses?", 1),
(
"Where is the plan for the vegetable bed and rainwater barrel?",
2,
),
] {
let query = embedder.embed_query(query)?;
let scores: Vec<f32> = vectors
.iter()
.map(|vector| vector.iter().zip(&query).map(|(a, b)| a * b).sum())
.collect();
let best = scores
.iter()
.enumerate()
.max_by(|a, b| a.1.total_cmp(b.1))
.unwrap()
.0;
assert_eq!(
best, expected,
"Semantic cross-language ranking failed: {scores:?}"
);
}
let long = format!("{}\nHotel accommodation on business trips is reimbursed after submitting the overnight stay receipt.",
"The routine maintenance log records ordinary uptime checks and cable inspections. ".repeat(350));
let windows = embedder.split_windows(&long, "maintenance-log.txt")?;
assert!(windows.len() > 1);
assert_eq!(windows.last().unwrap().end_byte, long.len());
assert!(windows
.last()
.unwrap()
.text
.contains("overnight stay receipt"));
let texts: Vec<_> = windows.iter().map(|window| window.text.as_str()).collect();
let vectors = embedder.embed_documents(&texts, "maintenance-log.txt")?;
let query = embedder.embed_query(
"Where does the policy say that lodging costs during work travel can be claimed?",
)?;
let scores: Vec<f32> = vectors
.iter()
.map(|vector| vector.iter().zip(&query).map(|(a, b)| a * b).sum())
.collect();
let best = scores
.iter()
.enumerate()
.max_by(|a, b| a.1.total_cmp(b.1))
.unwrap()
.0;
assert_eq!(
best,
windows.len() - 1,
"Tail retrieval did not find final window: {scores:?}"
);
Ok(())
}
crates/searchi-service/Cargo.toml
+20
View File
@@ -0,0 +1,20 @@
[package]
name = "searchi-service"
version.workspace = true
edition.workspace = true
license.workspace = true
description = "Persistent local hybrid retrieval over Searchi lexical and CPU semantic engines"
[dependencies]
anyhow.workspace = true
serde.workspace = true
searchi-core = { path = "../searchi-core" }
searchi-semantic = { path = "../searchi-semantic" }
rusqlite = { version = "0.32", features = ["bundled"] }
sha2 = "0.10"
fs2 = "0.4"
usearch = { version = "=2.26.4", default-features = false, features = ["numkong"] }
[dev-dependencies]
tempfile = "3"
serde_json.workspace = true
crates/searchi-service/benchmarks/windows-cpu-ann-100k.json
+16
View File
@@ -0,0 +1,16 @@
{
"ann_p95_ms": 1.3722,
"ann_recall_at10": 0.9850000000000001,
"database_ms": 39209.1035,
"dimensions": 768,
"graph_build_ms": 57950.0959,
"graph_memory_bytes": 554453184,
"hardware": "skylake",
"initial_eligibility_ms": 1612.2069999999999,
"kind": "synthetic index-only benchmark, not semantic-model quality",
"queries": 20,
"threads": 1,
"vectors": 100000,
"warm_passages_rehashed": 0,
"warm_service_p95_ms": 6.4033
}
crates/searchi-service/examples/evaluate.rs
+111
View File
@@ -0,0 +1,111 @@
//! Evaluate the immutable seed corpus with the actual configured Google model.
//! Usage: cargo run -p searchi-service --example evaluate -- MODEL_DIR WORKER OUTPUT_JSON
use anyhow::{ensure, Context, Result};
use searchi_core::SearchIndex;
use searchi_service::{HybridSearch, SearchOptions, SearchResult};
use serde_json::{json, Value};
use std::{fs, path::PathBuf, time::Instant};
fn expected_rank(results: &[SearchResult], filename: &str) -> Option<usize> {
results
.iter()
.position(|result| result.filename == filename)
.map(|rank| rank + 1)
}
fn main() -> Result<()> {
let args: Vec<_> = std::env::args_os().skip(1).collect();
ensure!(
args.len() == 3,
"Supply MODEL_DIR WORKER OUTPUT_JSON; this command never downloads assets"
);
ensure!(
std::env::var("SEARCHI_FORCE_CPU").is_ok_and(|value| value == "1"),
"Windows VM evaluation requires SEARCHI_FORCE_CPU=1"
);
let fixture: Value =
serde_json::from_str(include_str!("../../../fixtures/relevance/fixture.json"))?;
let temp = tempfile::tempdir()?;
let root = temp.path().join("documents");
fs::create_dir(&root)?;
let documents = fixture["documents"]
.as_array()
.context("documents missing")?;
for document in documents {
let filename = document["filename"].as_str().context("filename missing")?;
ensure!(
PathBuf::from(filename).file_name() == Some(std::ffi::OsStr::new(filename)),
"Fixture filename must be a basename"
);
fs::write(
root.join(filename),
document["text"].as_str().context("text missing")?,
)?;
}
let db = temp.path().join("index.db");
let mut lexical = SearchIndex::open(&db)?;
lexical.index_folder(&root)?;
let start = Instant::now();
let mut hybrid = HybridSearch::open(&db, &args[0], &args[1], 1)?;
let load_ms = start.elapsed().as_secs_f64() * 1000.0;
let start = Instant::now();
let indexed = hybrid.index_embeddings()?;
let indexing_ms = start.elapsed().as_secs_f64() * 1000.0;
drop(hybrid);
let start = Instant::now();
let mut hybrid = HybridSearch::open(&db, &args[0], &args[1], 1)?;
let startup_init_ms = start.elapsed().as_secs_f64() * 1000.0;
let start = Instant::now();
hybrid.warmup()?;
let background_prewarm_ms = start.elapsed().as_secs_f64() * 1000.0;
let start = Instant::now();
let first = hybrid.search(
"How can we end our relationship with the supplier?",
20,
&SearchOptions::default(),
)?;
let first_foreground_ms = start.elapsed().as_secs_f64() * 1000.0;
let first_foreground_top1 = first.first().map(|hit| hit.filename.as_str());
let mut rankings = Vec::new();
let mut durations = Vec::new();
for query in fixture["queries"].as_array().context("queries missing")? {
let text = query["text"].as_str().context("query text missing")?;
let expected = query["expected"].as_str().context("expected id missing")?;
let filename = documents
.iter()
.find(|doc| doc["id"] == expected)
.context("Expected document missing")?["filename"]
.as_str()
.unwrap();
let lexical_hits = lexical.search(text, 20)?;
let semantic_hits = hybrid.search_semantic(text, 20, &SearchOptions::default())?;
let start = Instant::now();
let hybrid_hits = hybrid.search(text, 20, &SearchOptions::default())?;
durations.push(start.elapsed().as_secs_f64() * 1000.0);
rankings.push(json!({"query":text,"expected":expected,"language":query["language"],
"lexical_rank":expected_rank(&lexical_hits,filename),
"semantic_rank":expected_rank(&semantic_hits,filename),
"hybrid_rank":expected_rank(&hybrid_hits,filename),
"hybrid_top3":hybrid_hits.iter().take(3).map(|hit| json!({"filename":hit.filename,"match_kind":hit.match_kind,"source":hit.source})).collect::<Vec<_>>()}));
}
durations.sort_by(f64::total_cmp);
let metrics = |lane: &str| {
json!({
"top1_accuracy":rankings.iter().filter(|row| row[lane]==1).count() as f64/rankings.len() as f64,
"recall_at3":rankings.iter().filter(|row| row[lane].as_u64().is_some_and(|rank| rank<=3)).count() as f64/rankings.len() as f64,
"recall_at10":rankings.iter().filter(|row| row[lane].as_u64().is_some_and(|rank| rank<=10)).count() as f64/rankings.len() as f64,
})
};
let report = json!({"model":"google/embeddinggemma-2","device":"cpu","threads":1,
"documents":documents.len(),"queries":rankings.len(),"load_ms":load_ms,"indexing_ms":indexing_ms,
"startup_init_ms":startup_init_ms,"background_prewarm_ms":background_prewarm_ms,"first_foreground_ms":first_foreground_ms,"first_foreground_top1":first_foreground_top1,
"hybrid_p50_ms":durations[durations.len()/2],"hybrid_p95_ms":durations[(durations.len()*95/100).min(durations.len()-1)],
"embedding_report":indexed,"lexical":metrics("lexical_rank"),"semantic":metrics("semantic_rank"),"hybrid":metrics("hybrid_rank"),
"rankings":rankings,"limitation":"Synthetic seed corpus only; no production quality or large-corpus latency certification"});
fs::write(&args[2], serde_json::to_string_pretty(&report)?)?;
println!(
"{}",
json!({"lexical":report["lexical"],"semantic":report["semantic"],"hybrid":report["hybrid"],"p95_ms":report["hybrid_p95_ms"]})
);
Ok(())
}
crates/searchi-service/src/ann_benchmark.rs
+163
View File
@@ -0,0 +1,163 @@
//! Explicit index benchmark: synthetic normalized vectors, not neural retrieval
//! certification. Includes the real generation store and warmed eligibility path.
use super::*;
use anyhow::{ensure, Result};
use std::{fs, time::Instant};
fn random(state: &mut u64) -> f32 {
*state ^= *state << 13;
*state ^= *state >> 7;
*state ^= *state << 17;
((*state >> 32) as u32 as f32 / u32::MAX as f32) * 2.0 - 1.0
}
#[test]
#[ignore = "Explicit CPU index benchmark; run cargo test --release -p searchi-service ann_100k -- --ignored --nocapture"]
fn ann_100k_quality_and_warm_metadata_latency() -> Result<()> {
ensure!(
!cfg!(debug_assertions),
"Run this performance benchmark with --release"
);
const COUNT: usize = 100_000;
const DIM: usize = 768;
const CLUSTERS: usize = 100;
let mut state = 0x19cfa024b9817731_u64;
let centers: Vec<Vec<f32>> = (0..CLUSTERS)
.map(|_| (0..DIM).map(|_| random(&mut state)).collect())
.collect();
let mut points = Vec::<Vec<f32>>::with_capacity(COUNT);
for i in 0..COUNT {
let mut vector: Vec<f32> = centers[i % CLUSTERS]
.iter()
.map(|value| value + 0.35 * random(&mut state))
.collect();
let norm = vector.iter().map(|v| v * v).sum::<f32>().sqrt();
for value in &mut vector {
*value /= norm;
}
points.push(vector);
}
let directory = tempfile::tempdir()?;
let db = directory.path().join("benchmark.db");
let lexical = searchi_core::SearchIndex::open(&db)?;
drop(lexical);
let mut store = Store::open(&db, DIM)?;
let generation = store.begin("synthetic-ann-benchmark-not-model-embeddings")?;
let started = Instant::now();
{
let transaction = store.connection.transaction()?;
let mut file=transaction.prepare("INSERT INTO files(path,root,filename,content,size,mtime,extension) VALUES(?1,?2,?3,?4,100,0,?5)")?;
let mut passage = transaction
.prepare("INSERT INTO passages(text,file_path,kind,number) VALUES(?1,?2,'line',1)")?;
let mut chunk=transaction.prepare("INSERT INTO semantic_chunks(generation,passage_id,ordinal,path,kind,number,content_hash,start_byte,end_byte,snippet,vector)
VALUES(?1,?2,0,?3,'line',1,?4,0,?5,?6,?7)")?;
for (i, vector) in points.iter().enumerate() {
let extension = if i % 10 == 0 { "md" } else { "txt" };
let filename = format!("document-{i}.{extension}");
let path = directory
.path()
.join(&filename)
.to_string_lossy()
.into_owned();
let text =
format!("Synthetic index benchmark record {i}. No inferred semantic labels.");
file.execute(params![
path,
directory.path().to_string_lossy(),
filename,
text,
extension
])?;
passage.execute(params![text, path])?;
let passage_id = transaction.last_insert_rowid();
let bytes: Vec<u8> = vector.iter().flat_map(|v| v.to_le_bytes()).collect();
chunk.execute(params![
generation,
passage_id,
path,
content_hash(&text),
text.len() as i64,
text,
bytes
])?;
}
drop(file);
drop(passage);
drop(chunk);
transaction.commit()?;
}
let database_ms = started.elapsed().as_secs_f64() * 1000.0;
let started = Instant::now();
store.activate(generation)?;
let graph_build_ms = started.elapsed().as_secs_f64() * 1000.0;
let ann = &store.ann_cache.as_ref().unwrap().1;
let hardware = ann.hardware_acceleration();
let graph_bytes = ann.memory_usage();
let query_ids: Vec<_> = (0..20).map(|i| (i * 4999 + 37) % COUNT).collect();
let mut recall = 0.0;
let mut ann_times = Vec::new();
for &query_id in &query_ids {
let query = &points[query_id];
let mut truth: Vec<_> = points
.iter()
.enumerate()
.map(|(i, vector)| (i, vector.iter().zip(query).map(|(a, b)| a * b).sum::<f32>()))
.collect();
truth.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
let expected: HashSet<_> = truth.iter().take(10).map(|(i, _)| *i as u64 + 1).collect();
let started = Instant::now();
let hits = ann.search(query, 10)?;
ann_times.push(started.elapsed().as_secs_f64() * 1000.0);
recall += hits
.keys
.iter()
.filter(|key| expected.contains(key))
.count() as f64
/ 10.0;
}
let started = Instant::now();
let _ = store.search(
&points[query_ids[0]],
"synthetic-ann-benchmark-not-model-embeddings",
10,
&SearchOptions::default(),
)?;
let initial_eligibility_ms = started.elapsed().as_secs_f64() * 1000.0;
let warmed_hashes = HASH_CALLS.with(std::cell::Cell::get);
let mut service_times = Vec::new();
for &query_id in &query_ids {
let started = Instant::now();
let _ = store.search(
&points[query_id],
"synthetic-ann-benchmark-not-model-embeddings",
10,
&SearchOptions::default(),
)?;
service_times.push(started.elapsed().as_secs_f64() * 1000.0);
}
ensure!(
HASH_CALLS.with(std::cell::Cell::get) == warmed_hashes,
"Warm benchmark queries rehashed unchanged passages"
);
ann_times.sort_by(f64::total_cmp);
service_times.sort_by(f64::total_cmp);
let p95 = |values: &[f64]| values[(values.len() * 95 / 100).min(values.len() - 1)];
let report = serde_json::json!({"kind":"synthetic index-only benchmark, not semantic-model quality",
"vectors":COUNT,"dimensions":DIM,"queries":query_ids.len(),"threads":1,"hardware":hardware,
"database_ms":database_ms,"graph_build_ms":graph_build_ms,"graph_memory_bytes":graph_bytes,
"ann_recall_at10":recall/query_ids.len() as f64,"ann_p95_ms":p95(&ann_times),
"initial_eligibility_ms":initial_eligibility_ms,"warm_service_p95_ms":p95(&service_times),"warm_passages_rehashed":0});
println!("{}", serde_json::to_string_pretty(&report)?);
if let Some(path) = std::env::var_os("SEARCHI_ANN_REPORT") {
fs::write(path, serde_json::to_string_pretty(&report)?)?;
}
ensure!(
report["ann_recall_at10"].as_f64().unwrap() >= 0.9,
"ANN recall quality gate failed; preserve report and tune graph parameters"
);
ensure!(
report["warm_service_p95_ms"].as_f64().unwrap() < 50.0,
"Warm retrieval latency gate failed"
);
Ok(())
}
crates/searchi-service/src/lib.rs
+531
View File
@@ -0,0 +1,531 @@
//! Shared offline hybrid search using real Google embeddings and persisted,
//! filtered HNSW retrieval. Production-scale validation remains a release gate.
mod store;
use anyhow::{ensure, Result};
use searchi_core::SearchIndex;
pub use searchi_core::{SearchOptions, SearchResult};
pub use searchi_semantic::WorkerInfo;
use searchi_semantic::{GemmaEmbedder, TextWindow, GEMMA_DIMENSIONS, GEMMA_FINGERPRINT};
use serde::Serialize;
use std::{collections::HashMap, path::Path};
use store::{Store, MAX_SCAN_CHUNKS};
const MAX_BATCH: usize = 8;
const BYTE_SLICE: usize = 64 * 1024;
const BYTE_OVERLAP: usize = 512;
const MAX_PASSAGE_WINDOWS: usize = 2048;
const CHUNK_POLICY: &str = "byte64k-overlap512-title512-v1";
fn generation_fingerprint(model: &str) -> String {
format!("{model}:{CHUNK_POLICY}")
}
/// Model-specific query/document prompts and actual tokenizer windows belong
/// to the provider. Storage and ranking never assume a particular dimension.
pub trait EmbeddingProvider: Send {
fn fingerprint(&self) -> &str;
fn dimensions(&self) -> usize;
fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>>;
fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result<Vec<Vec<f32>>>;
fn embed_query(&mut self, query: &str) -> Result<Vec<f32>>;
fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
None
}
}
impl EmbeddingProvider for GemmaEmbedder {
fn fingerprint(&self) -> &str {
GEMMA_FINGERPRINT
}
fn dimensions(&self) -> usize {
GEMMA_DIMENSIONS
}
fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>> {
self.split_windows(text, title)
}
fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result<Vec<Vec<f32>>> {
self.embed_documents(texts, title)
}
fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
self.embed_query(query)
}
fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
Some(self.info.clone())
}
}
#[derive(Debug, Serialize)]
pub struct EmbeddingReport {
pub passages: usize,
pub chunks: usize,
pub generation: i64,
pub model_fingerprint: String,
}
#[derive(Debug, Serialize)]
pub struct EmbeddingStatus {
pub generation: Option<i64>,
pub chunks: usize,
pub model_fingerprint: Option<String>,
pub stale: bool,
/// Rebuild is generation-based: all current passages are pending when stale.
pub pending: usize,
}
/// Shared local hybrid service. Google EmbeddingGemma 2 is the default and
/// only shipped provider; no alternative model is silently substituted.
pub struct HybridSearch {
lexical: SearchIndex,
embedder: Box<dyn EmbeddingProvider>,
store: Store,
generation_fingerprint: String,
}
impl HybridSearch {
/// Actual model/runtime/device reported by the persistent native worker.
pub fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
self.embedder.worker_info()
}
/// Cheap local status check. No model, runtime worker, or network is loaded.
/// Corpus revision triggers detect file inserts, updates, and deletions.
pub fn embedding_status(db: impl AsRef<Path>) -> Result<EmbeddingStatus> {
let connection =
rusqlite::Connection::open_with_flags(db, rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY)?;
connection.busy_timeout(std::time::Duration::from_secs(5))?;
let exists: bool = connection.query_row(
"SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE name='semantic_meta')",
[],
|row| row.get(0),
)?;
let pending: usize =
connection.query_row("SELECT count(*) FROM passages", [], |row| row.get(0))?;
if !exists {
return Ok(EmbeddingStatus {
generation: None,
chunks: 0,
model_fingerprint: None,
stale: true,
pending,
});
}
let version: i64 = connection.query_row(
"SELECT schema_version FROM semantic_meta WHERE singleton=1",
[],
|row| row.get(0),
)?;
if version < 5 {
return Ok(EmbeddingStatus {
generation: None,
chunks: 0,
model_fingerprint: None,
stale: true,
pending,
});
}
ensure!(
version == 5,
"Unsupported semantic schema version {version}"
);
let (generation,fingerprint,revision,current):(Option<i64>,Option<String>,Option<i64>,i64)=connection.query_row(
"SELECT g.id,g.fingerprint,CASE WHEN g.ann_blob IS NOT NULL AND g.dimensions=768 THEN g.corpus_revision END,m.corpus_revision FROM semantic_meta m LEFT JOIN semantic_generations g ON g.id=m.active_generation WHERE m.singleton=1",[],
|row| Ok((row.get(0)?,row.get(1)?,row.get(2)?,row.get(3)?)))?;
let chunks: usize = connection.query_row(
"SELECT count(*) FROM semantic_chunks WHERE generation=?1",
[generation],
|row| row.get(0),
)?;
let stale = revision != Some(current)
|| fingerprint.as_deref() != Some(generation_fingerprint(GEMMA_FINGERPRINT).as_str());
Ok(EmbeddingStatus {
generation,
chunks,
model_fingerprint: fingerprint,
stale,
pending: if stale { pending } else { 0 },
})
}
pub fn open(
db: impl AsRef<Path>,
model_dir: impl AsRef<Path>,
runtime_library: impl AsRef<Path>,
threads: usize,
) -> Result<Self> {
let embedder = GemmaEmbedder::load(model_dir, runtime_library, threads)?;
Self::with_provider(db, Box::new(embedder))
}
/// Explicit provider injection for future model generations, never fallback.
pub fn with_provider(
db: impl AsRef<Path>,
embedder: Box<dyn EmbeddingProvider>,
) -> Result<Self> {
let lexical = SearchIndex::open(&db)?;
let store = Store::open(db.as_ref(), embedder.dimensions())?;
let generation_fingerprint = generation_fingerprint(embedder.fingerprint());
Ok(Self {
lexical,
embedder,
store,
generation_fingerprint,
})
}
/// Execute real query inference and load the active ANN generation before
/// foreground use. Call on a background worker after assets are verified.
pub fn warmup(&mut self) -> Result<()> {
self.lexical.prewarm()?;
let vector = self.embedder.embed_query("warm up local file search")?;
let _ = self.store.search(
&vector,
&self.generation_fingerprint,
1,
&SearchOptions::default(),
)?;
Ok(())
}
/// Build a replacement generation in bounded CPU batches. Inference or
/// storage errors leave the previous generation active. A crashed build is
/// reclaimed on the next explicit rebuild under an OS advisory lock.
pub fn index_embeddings(&mut self) -> Result<EmbeddingReport> {
let _lock = self.store.build_lock()?;
let generation = self.store.begin(&self.generation_fingerprint)?;
let built = self.build(generation);
match built {
Ok(report) => {
self.store.activate(generation)?;
Ok(report)
}
Err(error) => {
self.store.abandon(generation)?;
Err(error)
}
}
}
fn build(&mut self, generation: i64) -> Result<EmbeddingReport> {
let mut report = EmbeddingReport {
passages: 0,
chunks: 0,
generation,
model_fingerprint: self.generation_fingerprint.clone(),
};
let mut after_id = 0;
loop {
let passages = self.lexical.passage_batch(after_id, 64)?;
if passages.is_empty() {
break;
}
for passage in passages {
after_id = passage.id;
if passage.text.trim().is_empty() {
continue;
}
let title = bounded_title(&passage.filename);
let windows = passage_windows(&mut *self.embedder, &passage.text, title)?;
ensure!(
report.chunks + windows.len() <= MAX_SCAN_CHUNKS,
"Corpus exceeds ANN generation resource budget of {MAX_SCAN_CHUNKS} chunks"
);
for (batch_index, batch) in windows.chunks(MAX_BATCH).enumerate() {
let texts: Vec<_> = batch.iter().map(|window| window.text.as_str()).collect();
let vectors = self.embedder.embed_documents(&texts, title)?;
self.store.insert(
generation,
&passage,
batch_index * MAX_BATCH,
batch,
&vectors,
)?;
}
report.chunks += windows.len();
report.passages += 1;
}
}
Ok(report)
}
/// Filter each retrieval lane before ranking, preserve filename priority,
/// then fuse keyword and semantic ranks with reciprocal-rank fusion (k=60).
pub fn search(
&mut self,
query: &str,
limit: usize,
options: &SearchOptions,
) -> Result<Vec<SearchResult>> {
let candidates = limit.saturating_mul(4).clamp(20, 1000);
// Core validates filters and query budgets even for empty requests.
let mut lexical = self
.lexical
.search_with_options(query, candidates, options)?;
if query.trim().is_empty() || limit == 0 {
return Ok(Vec::new());
}
promote_identifiers(query.trim(), &mut lexical);
let vector = self.embedder.embed_query(query.trim())?;
let semantic =
self.store
.search(&vector, &self.generation_fingerprint, candidates, options)?;
Ok(fuse(lexical, semantic, limit.min(1000)))
}
/// Semantic lane for reproducible relevance evaluation and diagnostics.
pub fn search_semantic(
&mut self,
query: &str,
limit: usize,
options: &SearchOptions,
) -> Result<Vec<SearchResult>> {
let _ = self.lexical.search_with_options(query, 0, options)?;
if query.trim().is_empty() || limit == 0 {
return Ok(Vec::new());
}
let vector = self.embedder.embed_query(query.trim())?;
self.store.search(
&vector,
&self.generation_fingerprint,
limit.min(1000),
options,
)
}
}
fn bounded_title(title: &str) -> &str {
let mut end = title.len().min(512);
while !title.is_char_boundary(end) {
end -= 1;
}
&title[..end]
}
fn byte_slices(text: &str) -> Result<Vec<(usize, usize)>> {
ensure!(
text.len() <= 2 * 1024 * 1024,
"Passage exceeds 2 MiB service byte budget"
);
let mut slices = Vec::new();
let mut start = 0;
while start < text.len() {
let mut end = (start + BYTE_SLICE).min(text.len());
while !text.is_char_boundary(end) {
end -= 1;
}
ensure!(end > start, "UTF-8 slice made no progress");
slices.push((start, end));
if end == text.len() {
break;
}
start = end.saturating_sub(BYTE_OVERLAP);
while !text.is_char_boundary(start) {
start -= 1;
}
}
Ok(slices)
}
fn passage_windows(
provider: &mut dyn EmbeddingProvider,
text: &str,
title: &str,
) -> Result<Vec<TextWindow>> {
let mut windows = Vec::new();
for (start, end) in byte_slices(text)? {
let fragment = &text[start..end];
if fragment.trim().is_empty() {
continue;
}
let mut slice = provider.split_windows(fragment, title)?;
for window in &mut slice {
window.start_byte += start;
window.end_byte += start;
}
windows.extend(slice);
ensure!(
windows.len() <= MAX_PASSAGE_WINDOWS,
"Passage exceeds {MAX_PASSAGE_WINDOWS} semantic window budget"
);
}
Ok(windows)
}
fn promote_identifiers(query: &str, results: &mut [SearchResult]) {
if query.chars().any(char::is_whitespace)
|| !query.chars().any(|c| c.is_ascii_digit() || c == '_')
{
return;
}
for result in results {
if result.match_kind == "content" && literal_identifier_match(&result.snippet, query) {
result.match_kind = "exact_identifier".into();
}
}
}
fn literal_identifier_match(text: &str, query: &str) -> bool {
text.match_indices(query).any(|(start, _)| {
let end = start + query.len();
let is_identifier = |c: char| c.is_alphanumeric() || c == '_' || c == '-';
!text[..start].chars().next_back().is_some_and(is_identifier)
&& !text[end..].chars().next().is_some_and(is_identifier)
})
}
fn fuse(
lexical: Vec<SearchResult>,
semantic: Vec<SearchResult>,
limit: usize,
) -> Vec<SearchResult> {
let mut results = HashMap::<String, SearchResult>::new();
for (rank, mut result) in lexical.into_iter().enumerate() {
let weight = if result.match_kind == "filename" {
1.15
} else {
1.0
};
result.score = weight / (61.0 + rank as f64);
results.insert(result.path.clone(), result);
}
for (rank, mut result) in semantic.into_iter().enumerate() {
let contribution = 1.0 / (61.0 + rank as f64);
if let Some(existing) = results.get_mut(&result.path) {
existing.score += contribution;
if existing.match_kind == "content" {
existing.match_kind = "hybrid".into();
}
} else {
result.score = contribution;
results.insert(result.path.clone(), result);
}
}
let tier = |result: &SearchResult| match result.match_kind.as_str() {
"exact_filename" => 0,
"exact_identifier" => 1,
_ => 2,
};
let mut results: Vec<_> = results.into_values().collect();
results.sort_by(|a, b| {
tier(a)
.cmp(&tier(b))
.then_with(|| b.score.total_cmp(&a.score))
.then_with(|| a.path.cmp(&b.path))
});
results.truncate(limit);
results
}
#[cfg(test)]
mod tests {
use super::*;
fn result(path: &str, kind: &str) -> SearchResult {
SearchResult {
path: path.into(),
filename: path.into(),
snippet: "evidence".into(),
score: 1.0,
match_kind: kind.into(),
source: None,
extraction_error: None,
}
}
#[test]
fn exact_filename_wins_even_when_another_result_matches_both_lanes() {
let results = fuse(
vec![
result("invoice.pdf", "exact_filename"),
result("notes.txt", "content"),
],
vec![
result("notes.txt", "semantic"),
result("another.txt", "semantic"),
],
10,
);
assert_eq!(results[0].path, "invoice.pdf");
assert_eq!(results[1].match_kind, "hybrid");
assert_eq!(results.len(), 3);
}
#[test]
fn identifiers_use_literal_boundaries_and_survive_semantic_competition() {
let mut lexical = vec![result("support.txt", "content")];
lexical[0].snippet = "Error E_SEARCH_071 means upgrade the app".into();
promote_identifiers("E_SEARCH_071", &mut lexical);
let results = fuse(lexical, vec![result("wrong.txt", "semantic")], 10);
assert_eq!(results[0].match_kind, "exact_identifier");
assert!(!literal_identifier_match(
"E_SEARCH_071_EXTRA",
"E_SEARCH_071"
));
}
#[test]
fn keyword_semantic_consensus_beats_weak_partial_filename() {
let results = fuse(
vec![
result("meeting-cancellation.txt", "filename"),
result("supplier-agreement.txt", "content"),
],
vec![result("supplier-agreement.txt", "semantic")],
10,
);
assert_eq!(results[0].path, "supplier-agreement.txt");
assert_eq!(results[0].match_kind, "hybrid");
assert_eq!(results[1].match_kind, "filename");
}
#[test]
fn byte_slices_preserve_utf8_overlap_offsets_and_complete_tail() {
let text = format!("{}TAIL", "Ä🐈中文 ".repeat(10000));
let slices = byte_slices(&text).unwrap();
assert!(slices.len() > 1);
assert_eq!(slices[0].0, 0);
assert_eq!(slices.last().unwrap().1, text.len());
for (start, end) in &slices {
assert!(text.is_char_boundary(*start) && text.is_char_boundary(*end));
assert!(end - start <= BYTE_SLICE);
}
for pair in slices.windows(2) {
assert!(pair[1].0 < pair[0].1 && pair[1].1 > pair[0].1);
}
assert!(text[slices.last().unwrap().0..].ends_with("TAIL"));
assert!(byte_slices(&"x".repeat(2 * 1024 * 1024 + 1)).is_err());
let title = "中文".repeat(120);
assert!(bounded_title(&title).len() <= 512 && title.starts_with(bounded_title(&title)));
}
#[test]
#[ignore = "Requires explicitly provisioned Google model and CPU worker"]
fn google_cpu_preserves_and_retrieves_content_after_64k_byte_boundary() -> Result<()> {
let model =
std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Set SEARCHI_GEMMA_MODEL_DIR");
let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Set SEARCHI_GEMMA_WORKER");
let mut provider =
GemmaEmbedder::load_with_device(model, worker, 1, searchi_semantic::Device::Cpu)?;
let text=format!("{}\nÜbernachtungskosten auf Dienstreisen werden gegen Vorlage der Hotelrechnung erstattet.","Routine system checks were completed successfully. ".repeat(1500));
assert!(text.len() > BYTE_SLICE);
let windows = passage_windows(&mut provider, &text, "routine-log.txt")?;
let last = windows.last().unwrap();
assert_eq!(last.end_byte, text.len());
assert!(last.start_byte >= BYTE_SLICE - BYTE_OVERLAP);
assert!(text.find("Übernachtungskosten").unwrap() > BYTE_SLICE);
assert!(last.text.contains("Hotelrechnung"));
assert_eq!(
text.get(last.start_byte..last.end_byte),
Some(last.text.as_str())
);
let vectors = provider.embed_documents(
&[windows[0].text.as_str(), last.text.as_str()],
"routine-log.txt",
)?;
let query = provider.embed_query(
"Where does it say that hotel stays during work travel can be reimbursed?",
)?;
let cosine = |vector: &[f32]| vector.iter().zip(&query).map(|(a, b)| a * b).sum::<f32>();
assert!(
cosine(&vectors[1]) > cosine(&vectors[0]),
"Real model did not retrieve the tail window"
);
Ok(())
}
}
crates/searchi-service/src/store.rs
+855
View File
@@ -0,0 +1,855 @@
use std::{
collections::{HashMap, HashSet},
fs::{File, OpenOptions},
path::{Path, PathBuf},
sync::Arc,
};
use anyhow::{ensure, Context, Result};
use fs2::FileExt;
use rusqlite::{params, Connection, OptionalExtension};
use searchi_core::{IndexedPassage, SearchOptions, SearchResult, SourceLocation};
use searchi_semantic::TextWindow;
use sha2::{Digest, Sha256};
use usearch::{Index, IndexOptions, MetricKind, ScalarKind};
pub(crate) const MAX_SCAN_CHUNKS: usize = 100_000;
pub(crate) const MIN_COSINE: f32 = 0.25;
#[derive(Clone, PartialEq, Eq)]
struct FilterKey {
folder: Option<String>,
extension: Option<String>,
min_size: Option<i64>,
max_size: Option<i64>,
modified_after: Option<i64>,
modified_before: Option<i64>,
}
struct EligibilityCache {
generation: i64,
revision: i64,
filter: FilterKey,
keys: Arc<HashSet<u64>>,
by_path: HashMap<String, Vec<u64>>,
}
pub(crate) struct Store {
connection: Connection,
lock_path: PathBuf,
dimensions: usize,
ann_cache: Option<(i64, Index)>,
eligibility_cache: Option<EligibilityCache>,
}
pub(crate) struct BuildLock(File);
impl Drop for BuildLock {
fn drop(&mut self) {
let _ = FileExt::unlock(&self.0);
}
}
pub(crate) fn content_hash(text: &str) -> String {
#[cfg(test)]
HASH_CALLS.with(|count| count.set(count.get() + 1));
format!("{:x}", Sha256::digest(text.as_bytes()))
}
#[cfg(test)]
thread_local! { static HASH_CALLS:std::cell::Cell<usize>=const {std::cell::Cell::new(0)}; }
impl Store {
pub(crate) fn open(db: &Path, dimensions: usize) -> Result<Self> {
ensure!(
(1..=4096).contains(&dimensions),
"Embedding dimensions must be between 1 and 4096"
);
ensure!(
db != Path::new(":memory:"),
"Hybrid search requires a persistent database path"
);
let mut connection = Connection::open(db)?;
connection.busy_timeout(std::time::Duration::from_secs(5))?;
connection.execute_batch("PRAGMA foreign_keys=ON;")?;
let initialization =
connection.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
initialization.execute_batch(
"CREATE TABLE IF NOT EXISTS semantic_meta (
singleton INTEGER PRIMARY KEY CHECK(singleton=1), schema_version INTEGER NOT NULL,
active_generation INTEGER, corpus_revision INTEGER NOT NULL DEFAULT 0, dirty_floor INTEGER NOT NULL DEFAULT 0);
INSERT OR IGNORE INTO semantic_meta(singleton,schema_version) VALUES(1,5);",
)?;
let version: i64 = initialization.query_row(
"SELECT schema_version FROM semantic_meta WHERE singleton=1",
[],
|row| row.get(0),
)?;
ensure!(
(1..=5).contains(&version),
"Unsupported semantic schema version {version}"
);
initialization.execute_batch("
CREATE TABLE IF NOT EXISTS semantic_generations(id INTEGER PRIMARY KEY, fingerprint TEXT NOT NULL, dimensions INTEGER NOT NULL, corpus_revision INTEGER NOT NULL DEFAULT 0, ann_blob BLOB, ann_sha256 TEXT);
CREATE TABLE IF NOT EXISTS semantic_chunks(
generation INTEGER NOT NULL REFERENCES semantic_generations(id) ON DELETE CASCADE,
passage_id INTEGER NOT NULL, ordinal INTEGER NOT NULL,
path TEXT NOT NULL, kind TEXT NOT NULL, number INTEGER NOT NULL,
content_hash TEXT NOT NULL, start_byte INTEGER NOT NULL, end_byte INTEGER NOT NULL,
snippet TEXT NOT NULL, vector BLOB NOT NULL,
PRIMARY KEY(generation,passage_id,ordinal));")?;
if version < 5 {
for (table,column,alter) in [
("semantic_meta","corpus_revision","ALTER TABLE semantic_meta ADD COLUMN corpus_revision INTEGER NOT NULL DEFAULT 0"),
("semantic_meta","dirty_floor","ALTER TABLE semantic_meta ADD COLUMN dirty_floor INTEGER NOT NULL DEFAULT 0"),
("semantic_generations","dimensions","ALTER TABLE semantic_generations ADD COLUMN dimensions INTEGER NOT NULL DEFAULT 384"),
("semantic_generations","corpus_revision","ALTER TABLE semantic_generations ADD COLUMN corpus_revision INTEGER NOT NULL DEFAULT 0"),
("semantic_generations","ann_blob","ALTER TABLE semantic_generations ADD COLUMN ann_blob BLOB"),
("semantic_generations","ann_sha256","ALTER TABLE semantic_generations ADD COLUMN ann_sha256 TEXT"),
] {
let exists:bool=initialization.query_row("SELECT EXISTS(SELECT 1 FROM pragma_table_info(?1) WHERE name=?2)",params![table,column],|row| row.get(0))?;
if !exists { initialization.execute_batch(alter)?; }
}
initialization.execute_batch("DROP TRIGGER IF EXISTS semantic_files_insert;
DROP TRIGGER IF EXISTS semantic_files_update; DROP TRIGGER IF EXISTS semantic_files_delete;")?;
}
initialization.execute_batch(
"CREATE TABLE IF NOT EXISTS semantic_dirty_paths(path TEXT PRIMARY KEY, revision INTEGER NOT NULL);
CREATE INDEX IF NOT EXISTS semantic_dirty_revision ON semantic_dirty_paths(revision);
CREATE INDEX IF NOT EXISTS semantic_chunk_path ON semantic_chunks(generation,path);
CREATE TRIGGER IF NOT EXISTS semantic_files_insert AFTER INSERT ON files BEGIN
UPDATE semantic_meta SET corpus_revision=corpus_revision+1 WHERE singleton=1;
INSERT INTO semantic_dirty_paths(path,revision) VALUES(new.path,(SELECT corpus_revision FROM semantic_meta WHERE singleton=1))
ON CONFLICT(path) DO UPDATE SET revision=excluded.revision; END;
CREATE TRIGGER IF NOT EXISTS semantic_files_update AFTER UPDATE ON files BEGIN
UPDATE semantic_meta SET corpus_revision=corpus_revision+1 WHERE singleton=1;
INSERT INTO semantic_dirty_paths(path,revision) VALUES(old.path,(SELECT corpus_revision FROM semantic_meta WHERE singleton=1))
ON CONFLICT(path) DO UPDATE SET revision=excluded.revision;
INSERT INTO semantic_dirty_paths(path,revision) VALUES(new.path,(SELECT corpus_revision FROM semantic_meta WHERE singleton=1))
ON CONFLICT(path) DO UPDATE SET revision=excluded.revision; END;
CREATE TRIGGER IF NOT EXISTS semantic_files_delete AFTER DELETE ON files BEGIN
UPDATE semantic_meta SET corpus_revision=corpus_revision+1 WHERE singleton=1;
INSERT INTO semantic_dirty_paths(path,revision) VALUES(old.path,(SELECT corpus_revision FROM semantic_meta WHERE singleton=1))
ON CONFLICT(path) DO UPDATE SET revision=excluded.revision; END;
UPDATE semantic_meta SET schema_version=5 WHERE singleton=1;",
)?;
initialization.commit()?;
let canonical = db.canonicalize()?;
let lock_path = PathBuf::from(format!("{}.semantic.lock", canonical.to_string_lossy()));
Ok(Self {
connection,
lock_path,
dimensions,
ann_cache: None,
eligibility_cache: None,
})
}
pub(crate) fn build_lock(&self) -> Result<BuildLock> {
let file = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(&self.lock_path)?;
file.try_lock_exclusive()
.context("Another semantic rebuild is running")?;
Ok(BuildLock(file))
}
pub(crate) fn begin(&mut self, fingerprint: &str) -> Result<i64> {
let transaction = self.connection.transaction()?;
// Only called while holding the cross-process lock: abandoned staging
// generations may be removed safely after a crash.
transaction.execute("DELETE FROM semantic_generations WHERE id != COALESCE((SELECT active_generation FROM semantic_meta WHERE singleton=1),-1)", [])?;
transaction.execute(
"INSERT INTO semantic_generations(fingerprint,dimensions,corpus_revision) SELECT ?1,?2,corpus_revision FROM semantic_meta WHERE singleton=1",
params![fingerprint, self.dimensions as i64],
)?;
let id = transaction.last_insert_rowid();
transaction.commit()?;
Ok(id)
}
pub(crate) fn insert(
&mut self,
generation: i64,
passage: &IndexedPassage,
first_ordinal: usize,
windows: &[TextWindow],
vectors: &[Vec<f32>],
) -> Result<()> {
ensure!(
windows.len() == vectors.len(),
"Window and embedding counts differ"
);
let hash = content_hash(&passage.text);
let transaction = self.connection.transaction()?;
for (i, (window, vector)) in windows.iter().zip(vectors).enumerate() {
ensure!(
vector.len() == self.dimensions && vector.iter().all(|v| v.is_finite()),
"Invalid persisted vector"
);
let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
ensure!(
norm.is_finite() && (norm - 1.0).abs() < 1e-3,
"Persisted embeddings must be L2 normalized"
);
ensure!(
passage.text.get(window.start_byte..window.end_byte) == Some(window.text.as_str()),
"Window does not match passage source range"
);
let bytes: Vec<u8> = vector
.iter()
.flat_map(|value| value.to_le_bytes())
.collect();
transaction.execute("INSERT INTO semantic_chunks(generation,passage_id,ordinal,path,kind,number,content_hash,start_byte,end_byte,snippet,vector)
VALUES(?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11)", params![generation, passage.id,
(first_ordinal+i) as i64, passage.path, passage.source.kind, passage.source.number,
hash, window.start_byte as i64, window.end_byte as i64, window.text, bytes])?;
}
transaction.commit()?;
Ok(())
}
pub(crate) fn activate(&mut self, generation: i64) -> Result<()> {
// A failed graph build or serialization leaves the old generation active.
let ann = self.build_ann(generation)?;
let length = ann.serialized_length();
ensure!(
length <= 512 * 1024 * 1024,
"Serialized ANN exceeds 512 MiB generation budget"
);
let mut serialized = vec![0_u8; length];
ann.save_to_buffer(&mut serialized)?;
let checksum = format!("{:x}", Sha256::digest(&serialized));
let transaction = self.connection.transaction()?;
ensure!(
transaction.query_row(
"SELECT EXISTS(SELECT 1 FROM semantic_generations WHERE id=?1)",
[generation],
|row| row.get::<_, bool>(0)
)?,
"Missing staging generation"
);
transaction.execute(
"UPDATE semantic_generations SET ann_blob=?1,ann_sha256=?2 WHERE id=?3",
params![serialized, checksum, generation],
)?;
transaction.execute(
"UPDATE semantic_meta SET active_generation=?1 WHERE singleton=1",
[generation],
)?;
transaction.execute(
"DELETE FROM semantic_generations WHERE id!=?1",
[generation],
)?;
transaction.execute_batch("UPDATE semantic_meta SET dirty_floor=max(dirty_floor,corpus_revision-100000) WHERE singleton=1;
DELETE FROM semantic_dirty_paths WHERE revision<(SELECT dirty_floor FROM semantic_meta WHERE singleton=1);")?;
transaction.commit()?;
self.ann_cache = Some((generation, ann));
self.eligibility_cache = None;
Ok(())
}
fn new_ann(&self) -> Result<Index> {
Ok(Index::new(&IndexOptions {
dimensions: self.dimensions,
metric: MetricKind::Cos,
quantization: ScalarKind::F32,
connectivity: 16,
expansion_add: 128,
expansion_search: 64,
multi: false,
})?)
}
fn build_ann(&self, generation: i64) -> Result<Index> {
let count: usize = self.connection.query_row(
"SELECT count(*) FROM semantic_chunks WHERE generation=?1",
[generation],
|row| row.get(0),
)?;
ensure!(
count <= MAX_SCAN_CHUNKS,
"ANN generation exceeds {MAX_SCAN_CHUNKS} chunks"
);
let ann = self.new_ann()?;
ann.reserve_capacity_and_threads(count.max(1), 1)?;
let mut statement = self.connection.prepare(
"SELECT rowid,vector FROM semantic_chunks WHERE generation=?1 ORDER BY rowid",
)?;
let mut rows = statement.query([generation])?;
while let Some(row) = rows.next()? {
let key: i64 = row.get(0)?;
let bytes: Vec<u8> = row.get(1)?;
ensure!(
bytes.len() == self.dimensions * 4,
"Invalid staging embedding length"
);
let vector: Vec<f32> = bytes
.as_chunks::<4>()
.0
.iter()
.map(|bytes| f32::from_le_bytes(*bytes))
.collect();
ensure!(
vector.iter().all(|v| v.is_finite()),
"Invalid staging embedding values"
);
ann.add(u64::try_from(key)?, &vector)?;
}
Ok(ann)
}
pub(crate) fn abandon(&mut self, generation: i64) -> Result<()> {
self.connection.execute("DELETE FROM semantic_generations WHERE id=?1 AND id!=COALESCE((SELECT active_generation FROM semantic_meta WHERE singleton=1),-1)", [generation])?;
Ok(())
}
pub(crate) fn search(
&mut self,
query: &[f32],
fingerprint: &str,
limit: usize,
options: &SearchOptions,
) -> Result<Vec<SearchResult>> {
ensure!(
query.len() == self.dimensions && query.iter().all(|v| v.is_finite()),
"Invalid query vector"
);
if limit == 0 {
return Ok(Vec::new());
}
// A consistent read snapshot covers eligibility and materialization.
let snapshot = self.connection.unchecked_transaction()?;
let active:Option<i64>=snapshot.query_row("SELECT g.id FROM semantic_generations g JOIN semantic_meta m ON m.active_generation=g.id WHERE g.fingerprint=?1 AND g.dimensions=?2",
params![fingerprint,self.dimensions as i64],|row|row.get(0)).optional()?;
let Some(active) = active else {
return Ok(Vec::new());
};
if !self
.ann_cache
.as_ref()
.is_some_and(|(generation, _)| *generation == active)
{
let length: Option<usize> = snapshot.query_row(
"SELECT length(ann_blob) FROM semantic_generations WHERE id=?1",
[active],
|row| row.get(0),
)?;
ensure!(
length.is_some_and(|bytes| bytes <= 512 * 1024 * 1024),
"Missing or oversized ANN generation; rebuild embeddings explicitly"
);
let (serialized, checksum): (Vec<u8>, String) = snapshot.query_row(
"SELECT ann_blob,ann_sha256 FROM semantic_generations WHERE id=?1",
[active],
|row| Ok((row.get(0)?, row.get(1)?)),
)?;
ensure!(
format!("{:x}", Sha256::digest(&serialized)) == checksum,
"ANN generation SHA256 mismatch; rebuild embeddings explicitly"
);
let ann = self.new_ann()?;
ann.load_from_buffer(&serialized)?;
ensure!(
ann.dimensions() == self.dimensions,
"ANN dimensions differ from embedding provider"
);
ann.reserve_capacity_and_threads(ann.size().max(1), 1)?;
self.ann_cache = Some((active, ann));
}
let folder = options
.folder
.as_ref()
.map(std::fs::canonicalize)
.transpose()?
.map(|p| {
format!(
"{}{}",
p.to_string_lossy().trim_end_matches(['/', '\\']),
std::path::MAIN_SEPARATOR
)
});
let extension = options
.extension
.as_ref()
.map(|ext| ext.trim_start_matches('.').to_ascii_lowercase());
let min_size = options.min_size.map(i64::try_from).transpose()?;
let max_size = options.max_size.map(i64::try_from).transpose()?;
let filter = FilterKey {
folder,
extension,
min_size,
max_size,
modified_after: options.modified_after_unix_ms,
modified_before: options.modified_before_unix_ms,
};
let (revision, floor): (i64, i64) = snapshot.query_row(
"SELECT corpus_revision,dirty_floor FROM semantic_meta WHERE singleton=1",
[],
|row| Ok((row.get(0)?, row.get(1)?)),
)?;
let existing = self.eligibility_cache.take();
let mut cache = match existing {
Some(mut cache)
if cache.generation == active
&& cache.filter == filter
&& cache.revision >= floor =>
{
if cache.revision != revision {
let mut statement=snapshot.prepare("SELECT path FROM semantic_dirty_paths WHERE revision>?1 AND revision<=?2 ORDER BY path LIMIT 1001")?;
let paths = statement
.query_map(params![cache.revision, revision], |row| {
row.get::<_, String>(0)
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
if paths.len() > 1000 {
cache = collect_eligibility(&snapshot, active, revision, &filter, None)?;
} else {
for path in paths {
if let Some(old) = cache.by_path.remove(&path) {
for key in old {
Arc::make_mut(&mut cache.keys).remove(&key);
}
}
let changed = collect_eligibility(
&snapshot,
active,
revision,
&filter,
Some(&path),
)?;
Arc::make_mut(&mut cache.keys).extend(changed.keys.iter().copied());
cache.by_path.extend(changed.by_path);
}
cache.revision = revision;
}
}
cache
}
_ => collect_eligibility(&snapshot, active, revision, &filter, None)?,
};
ensure!(
cache.keys.len() <= MAX_SCAN_CHUNKS,
"Eligibility cache exceeds chunk resource budget"
);
cache.revision = revision;
let eligible = Arc::clone(&cache.keys);
self.eligibility_cache = Some(cache);
if eligible.is_empty() {
return Ok(Vec::new());
}
let ann = &self.ann_cache.as_ref().expect("loaded ANN cache").1;
let mut candidate_count = limit.saturating_mul(4).min(eligible.len()).max(1);
let mut results;
loop {
let matches =
ann.filtered_search(query, candidate_count, |key| eligible.contains(&key))?;
let mut ranked: Vec<_> = matches.keys.into_iter().zip(matches.distances).collect();
ranked.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
let mut best = HashMap::<String, SearchResult>::new();
let mut materialize=snapshot.prepare("SELECT c.path,f.filename,c.snippet,c.kind,c.number FROM semantic_chunks c JOIN files f ON f.path=c.path WHERE c.rowid=?1 AND c.generation=?2")?;
for (key, distance) in ranked {
let cosine = 1.0 - distance;
ensure!(cosine.is_finite(), "ANN produced nonfinite distance");
if cosine < MIN_COSINE {
continue;
}
let result =
materialize.query_row(params![i64::try_from(key)?, active], |row| {
Ok(SearchResult {
path: row.get(0)?,
filename: row.get(1)?,
snippet: row.get(2)?,
score: f64::from(cosine),
match_kind: "semantic".into(),
source: Some(SourceLocation {
kind: row.get(3)?,
number: row.get(4)?,
}),
extraction_error: None,
})
})?;
if best.contains_key(&result.path) {
continue;
}
best.insert(result.path.clone(), result);
}
results = best.into_values().collect::<Vec<_>>();
if results.len() >= limit || candidate_count == eligible.len() {
break;
}
candidate_count = candidate_count.saturating_mul(2).min(eligible.len());
}
results.sort_by(|a, b| {
b.score
.total_cmp(&a.score)
.then_with(|| a.path.cmp(&b.path))
});
results.truncate(limit);
snapshot.commit()?;
Ok(results)
}
#[cfg(test)]
fn exact_search(
&self,
query: &[f32],
fingerprint: &str,
limit: usize,
options: &SearchOptions,
) -> Result<Vec<SearchResult>> {
ensure!(
query.len() == self.dimensions && query.iter().all(|v| v.is_finite()),
"Invalid query vector"
);
let active: Option<i64> = self.connection.query_row("SELECT g.id FROM semantic_generations g JOIN semantic_meta m ON m.active_generation=g.id WHERE g.fingerprint=?1 AND g.dimensions=?2", params![fingerprint,self.dimensions as i64], |row| row.get(0)).optional()?;
let Some(active) = active else {
return Ok(Vec::new());
};
let folder = options
.folder
.as_ref()
.map(std::fs::canonicalize)
.transpose()?
.map(|p| {
format!(
"{}{}",
p.to_string_lossy().trim_end_matches(['/', '\\']),
std::path::MAIN_SEPARATOR
)
});
let extension = options
.extension
.as_ref()
.map(|ext| ext.trim_start_matches('.').to_ascii_lowercase());
let min_size = options.min_size.map(i64::try_from).transpose()?;
let max_size = options.max_size.map(i64::try_from).transpose()?;
let mut statement = self.connection.prepare("SELECT c.passage_id,c.path,f.filename,c.snippet,c.vector,c.content_hash,p.text,c.kind,c.number
FROM semantic_chunks c JOIN passages p ON p.rowid=c.passage_id AND p.file_path=c.path AND p.kind=c.kind AND p.number=c.number
JOIN files f ON f.path=c.path WHERE c.generation=?1
AND (?2 IS NULL OR substr(f.path,1,length(?2))=?2)
AND (?3 IS NULL OR f.extension=?3)
AND (?4 IS NULL OR f.size>=?4) AND (?5 IS NULL OR f.size<=?5)
AND (?6 IS NULL OR f.mtime/1000000>=?6) AND (?7 IS NULL OR f.mtime/1000000<=?7)
ORDER BY c.passage_id,c.ordinal LIMIT ?8")?;
// Fetch one extra to detect the baseline cap instead of silently hiding
// candidates. Metadata filtering precedes this cap and all top-k limits.
let mut rows = statement.query(params![
active,
folder,
extension,
min_size,
max_size,
options.modified_after_unix_ms,
options.modified_before_unix_ms,
MAX_SCAN_CHUNKS as i64 + 1
])?;
let mut hashes = HashMap::<i64, String>::new();
let mut best = HashMap::<String, SearchResult>::new();
let mut scanned = 0;
while let Some(row) = rows.next()? {
scanned += 1;
ensure!(
scanned <= MAX_SCAN_CHUNKS,
"Exact-vector baseline exceeds {MAX_SCAN_CHUNKS} chunks; ANN integration required"
);
let passage_id: i64 = row.get(0)?;
let current_hash = if let Some(hash) = hashes.get(&passage_id) {
hash.clone()
} else {
let hash = content_hash(&row.get::<_, String>(6)?);
hashes.insert(passage_id, hash.clone());
hash
};
if current_hash != row.get::<_, String>(5)? {
continue;
}
let bytes: Vec<u8> = row.get(4)?;
ensure!(
bytes.len() == self.dimensions * 4,
"Corrupt stored embedding length"
);
let mut cosine = 0.0;
for (dimension, value) in bytes.as_chunks::<4>().0.iter().enumerate() {
let value = f32::from_le_bytes(*value);
ensure!(value.is_finite(), "Corrupt stored embedding value");
cosine += value * query[dimension];
}
if cosine < MIN_COSINE {
continue;
}
let path: String = row.get(1)?;
if best
.get(&path)
.is_some_and(|result| result.score >= f64::from(cosine))
{
continue;
}
best.insert(
path.clone(),
SearchResult {
path,
filename: row.get(2)?,
snippet: row.get(3)?,
score: f64::from(cosine),
match_kind: "semantic".into(),
source: Some(SourceLocation {
kind: row.get(7)?,
number: row.get(8)?,
}),
extraction_error: None,
},
);
}
let mut results: Vec<_> = best.into_values().collect();
results.sort_by(|a, b| {
b.score
.total_cmp(&a.score)
.then_with(|| a.path.cmp(&b.path))
});
results.truncate(limit);
Ok(results)
}
}
fn collect_eligibility(
connection: &Connection,
generation: i64,
revision: i64,
filter: &FilterKey,
path: Option<&str>,
) -> Result<EligibilityCache> {
let mut statement=connection.prepare("SELECT c.rowid,c.passage_id,c.content_hash,p.text,c.path FROM semantic_chunks c
JOIN passages p ON p.rowid=c.passage_id AND p.file_path=c.path AND p.kind=c.kind AND p.number=c.number
JOIN files f ON f.path=c.path WHERE c.generation=?1
AND (?2 IS NULL OR substr(f.path,1,length(?2))=?2) AND (?3 IS NULL OR f.extension=?3)
AND (?4 IS NULL OR f.size>=?4) AND (?5 IS NULL OR f.size<=?5)
AND (?6 IS NULL OR f.mtime/1000000>=?6) AND (?7 IS NULL OR f.mtime/1000000<=?7)
AND (?9 IS NULL OR c.path=?9) ORDER BY c.passage_id,c.ordinal LIMIT ?8")?;
let mut rows = statement.query(params![
generation,
filter.folder,
filter.extension,
filter.min_size,
filter.max_size,
filter.modified_after,
filter.modified_before,
MAX_SCAN_CHUNKS as i64 + 1,
path
])?;
let mut hashes = HashMap::<i64, String>::new();
let mut keys = HashSet::new();
let mut by_path = HashMap::<String, Vec<u64>>::new();
let mut count = 0;
while let Some(row) = rows.next()? {
count += 1;
ensure!(
count <= MAX_SCAN_CHUNKS,
"Corpus exceeds {MAX_SCAN_CHUNKS} chunk resource budget"
);
let passage_id: i64 = row.get(1)?;
let hash = if let Some(hash) = hashes.get(&passage_id) {
hash.clone()
} else {
let hash = content_hash(&row.get::<_, String>(3)?);
hashes.insert(passage_id, hash.clone());
hash
};
if hash == row.get::<_, String>(2)? {
let key = u64::try_from(row.get::<_, i64>(0)?)?;
keys.insert(key);
by_path.entry(row.get(4)?).or_default().push(key);
}
}
Ok(EligibilityCache {
generation,
revision,
filter: filter.clone(),
keys: Arc::new(keys),
by_path,
})
}
#[cfg(test)]
mod tests {
use super::*;
use searchi_core::SearchIndex;
const DIMENSIONS: usize = 768;
#[test]
fn generation_persists_failed_staging_preserves_active_and_changes_hide_stale_vectors(
) -> Result<()> {
let directory = tempfile::tempdir()?;
let root = directory.path().join("docs");
std::fs::create_dir(&root)?;
let db = directory.path().join("index.db");
std::fs::write(root.join("notes.txt"), "original relevant content")?;
let mut index = SearchIndex::open(&db)?;
index.index_folder(&root)?;
let passage = index.passage_batch(0, 1)?.remove(0);
let mut vector = vec![0.0; DIMENSIONS];
vector[0] = 1.0;
let windows = [TextWindow {
text: passage.text.clone(),
start_byte: 0,
end_byte: passage.text.len(),
}];
let mut store = Store::open(&db, DIMENSIONS)?;
let generation = store.begin("test-storage-vectors")?;
store.insert(generation, &passage, 0, &windows, &[vector.clone()])?;
store.activate(generation)?;
assert_eq!(
store.exact_search(
&vector,
"test-storage-vectors",
10,
&SearchOptions::default()
)?[0]
.path,
store.search(
&vector,
"test-storage-vectors",
10,
&SearchOptions::default()
)?[0]
.path
);
drop(store);
let mut store = Store::open(&db, DIMENSIONS)?;
assert_eq!(
store
.search(
&vector,
"test-storage-vectors",
10,
&SearchOptions::default()
)?
.len(),
1
);
assert!(store
.search(&vector, "different-model", 10, &SearchOptions::default())?
.is_empty());
let staging = store.begin("new-model")?;
assert_eq!(
store
.search(
&vector,
"test-storage-vectors",
10,
&SearchOptions::default()
)?
.len(),
1
);
store.abandon(staging)?;
std::fs::write(
root.join("notes.txt"),
"replacement unrelated content longer than before",
)?;
index.index_folder(&root)?;
assert!(store
.search(
&vector,
"test-storage-vectors",
10,
&SearchOptions::default()
)?
.is_empty());
std::fs::remove_file(root.join("notes.txt"))?;
index.index_folder(&root)?;
assert!(store
.search(
&vector,
"test-storage-vectors",
10,
&SearchOptions::default()
)?
.is_empty());
Ok(())
}
#[test]
fn ann_filters_before_top_k_checks_checksum_and_status_tracks_corpus_changes() -> Result<()> {
let directory = tempfile::tempdir()?;
let root = directory.path().join("docs");
std::fs::create_dir(&root)?;
let db = directory.path().join("index.db");
std::fs::write(root.join("closest.md"), "closest metadata excluded source")?;
std::fs::write(root.join("allowed.txt"), "allowed metadata source")?;
let mut lexical = SearchIndex::open(&db)?;
lexical.index_folder(&root)?;
let passages = lexical.passage_batch(0, 10)?;
let mut store = Store::open(&db, DIMENSIONS)?;
let fingerprint = crate::generation_fingerprint(searchi_semantic::GEMMA_FINGERPRINT);
let generation = store.begin(&fingerprint)?;
let mut query = vec![0.0; DIMENSIONS];
query[0] = 1.0;
for passage in &passages {
let mut vector = query.clone();
if passage.filename == "allowed.txt" {
vector[0] = 0.8;
vector[1] = 0.6;
}
let windows = [TextWindow {
text: passage.text.clone(),
start_byte: 0,
end_byte: passage.text.len(),
}];
store.insert(generation, passage, 0, &windows, &[vector])?;
}
store.activate(generation)?;
let options = SearchOptions {
extension: Some("txt".into()),
..Default::default()
};
assert_eq!(
store.search(&query, &fingerprint, 1, &options)?[0].filename,
"allowed.txt"
);
let warmed_hashes = HASH_CALLS.with(std::cell::Cell::get);
let _ = store.search(&query, &fingerprint, 1, &options)?;
assert_eq!(
HASH_CALLS.with(std::cell::Cell::get),
warmed_hashes,
"Warm query rehashed unchanged corpus"
);
assert_eq!(
store.exact_search(&query, &fingerprint, 1, &options)?[0].filename,
"allowed.txt"
);
assert!(!crate::HybridSearch::embedding_status(&db)?.stale);
std::fs::write(
root.join("allowed.txt"),
"replacement changes the generation revision",
)?;
lexical.index_folder(&root)?;
let status = crate::HybridSearch::embedding_status(&db)?;
let before_delta = HASH_CALLS.with(std::cell::Cell::get);
assert!(store.search(&query, &fingerprint, 1, &options)?.is_empty());
assert!(
HASH_CALLS.with(std::cell::Cell::get) - before_delta <= 1,
"Single file change rehashed unrelated passages"
);
assert!(status.stale);
assert_eq!(status.pending, 2);
store.connection.execute(
"UPDATE semantic_generations SET ann_blob=x'010203' WHERE id=?1",
[generation],
)?;
drop(store);
let mut store = Store::open(&db, DIMENSIONS)?;
assert!(store
.search(&query, &fingerprint, 1, &SearchOptions::default())
.unwrap_err()
.to_string()
.contains("SHA256 mismatch"));
Ok(())
}
}
#[cfg(test)]
#[path = "ann_benchmark.rs"]
mod ann_benchmark;
crates/searchi-service/tests/google_cpu.rs
+80
View File
@@ -0,0 +1,80 @@
use anyhow::Result;
use searchi_core::SearchIndex;
use searchi_service::{HybridSearch, SearchOptions};
use std::fs;
#[test]
#[ignore = "Requires explicitly provisioned Google model and CPU worker"]
fn real_google_search_persists_and_suppresses_changed_deleted_sources() -> Result<()> {
assert_eq!(
std::env::var("SEARCHI_FORCE_CPU").unwrap(),
"1",
"VM tests require CPU only"
);
let model = std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Set SEARCHI_GEMMA_MODEL_DIR");
let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Set SEARCHI_GEMMA_WORKER");
let temp = tempfile::tempdir()?;
let root = temp.path().join("documents");
fs::create_dir(&root)?;
let db = temp.path().join("index.db");
fs::write(
root.join("agreement.txt"),
"The supplier partnership may be ended with sixty days written notice.",
)?;
fs::write(
root.join("garden.txt"),
"Tomatoes need regular watering and lots of sunlight.",
)?;
let mut lexical = SearchIndex::open(&db)?;
lexical.index_folder(&root)?;
let mut hybrid = HybridSearch::open(&db, &model, &worker, 1)?;
let report = hybrid.index_embeddings()?;
assert_eq!(report.passages, 2);
let hits = hybrid.search(
"Wie beenden wir die Zusammenarbeit mit dem Lieferanten?",
10,
&SearchOptions::default(),
)?;
assert_eq!(hits[0].filename, "agreement.txt");
assert!(hits[0].source.is_some());
assert!(hits[0].snippet.contains("sixty days"));
let exact = hybrid.search("garden.txt", 10, &SearchOptions::default())?;
assert_eq!(exact[0].match_kind, "exact_filename");
let restricted = hybrid.search(
"Wie beenden wir die Zusammenarbeit mit dem Lieferanten?",
10,
&SearchOptions {
extension: Some("pdf".into()),
..Default::default()
},
)?;
assert!(restricted.is_empty());
drop(hybrid);
let mut hybrid = HybridSearch::open(&db, &model, &worker, 1)?;
assert_eq!(
hybrid.search(
"How can we stop working with the vendor?",
10,
&SearchOptions::default()
)?[0]
.filename,
"agreement.txt"
);
fs::write(root.join("agreement.txt"),"Entirely unrelated replacement content about repainting a kitchen and changing cupboard handles.")?;
lexical.index_folder(&root)?;
assert!(hybrid
.search_semantic(
"How can we stop working with the vendor?",
10,
&SearchOptions::default()
)?
.iter()
.all(|hit| hit.filename != "agreement.txt"));
fs::remove_file(root.join("garden.txt"))?;
lexical.index_folder(&root)?;
assert!(hybrid
.search_semantic("How do we water vegetables?", 10, &SearchOptions::default())?
.iter()
.all(|hit| hit.filename != "garden.txt"));
Ok(())
}