feat: add desktop document search and verified native CPU embeddings
Core CI / core (ubuntu-latest) (push) Failing after 1m22s
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 3m46s
Core CI / desktop (macos-latest) (push) Canceled after 0s
Core CI / desktop (windows-latest) (push) Canceled after 0s

This commit is contained in:
karl committed 2026-10-12 16:07:53 +02:00
1 parent f77112db97
commit a9b6592330
46 files changed
+10794 -77

No files matched your search

+11 -3
View File
@@ -7,6 +7,8 @@ GPU setup. Index only folders you explicitly choose.
cargo run -p searchi-cli -- index --db ./searchi.db --root ./documents
cargo run -p searchi-cli -- search --db ./searchi.db "termination clause" --limit 10
cargo run -p searchi-cli -- search --db ./searchi.db "invoice" --json
cargo run -p searchi-cli -- search --db ./searchi.db "invoice" --extension pdf --min-size 1024 --json
cargo run -p searchi-cli -- index --db ./searchi.db --root ./documents --exclude private
cargo run -p searchi-cli -- stats --db ./searchi.db --json
```
@@ -23,6 +25,12 @@ snippet, ranking score, and match kind. An empty search emits `[]`. `stats --jso
emits an object with `files` and `content_files`. JSON is the scripting interface;
the human-readable output may change as preview support develops.
The foundation CLI searches filenames and extracted text. Semantic inference,
document formats beyond the core extractor, and desktop integration remain
tracked in the project roadmap; this interface does not fabricate embeddings.
The CLI searches filenames, bounded text, PDF text and DOCX paragraphs. Results
include actual matching source locations when available and extraction errors
for documents that could not be read. Semantic inference and media formats
remain tracked in the roadmap; this interface does not fabricate embeddings.
Filters apply before result limits: `--folder`, `--extension`, `--min-size`,
`--max-size`, `--modified-after-ms`, and `--modified-before-ms`. Dates are Unix
milliseconds and inclusive. Repeated `index --exclude NAME` flags add directory
names to the defaults (`.git`, `node_modules`, `target`).
+49 -6
View File
@@ -4,7 +4,7 @@ use std::path::PathBuf;
use anyhow::{bail, Context, Result};
use clap::{Parser, Subcommand};
use searchi_core::SearchIndex;
use searchi_core::{IndexOptions, SearchIndex, SearchOptions};
#[derive(Debug, Parser)]
#[command(name = "searchi", version, about = "Search your local files offline")]
@@ -23,6 +23,9 @@ enum Command {
/// Folder to index recursively.
#[arg(long)]
root: PathBuf,
/// Additional directory names to exclude; repeat this flag as needed.
#[arg(long)]
exclude: Vec<String>,
},
/// Search filenames and indexed document content.
Search {
@@ -36,6 +39,21 @@ enum Command {
/// Emit a JSON array for scripting.
#[arg(long)]
json: bool,
/// Restrict results to this folder and descendants.
#[arg(long)]
folder: Option<PathBuf>,
/// File extension, with or without the leading dot.
#[arg(long)]
extension: Option<String>,
#[arg(long)]
min_size: Option<u64>,
#[arg(long)]
max_size: Option<u64>,
/// Inclusive modification time in Unix milliseconds.
#[arg(long)]
modified_after_ms: Option<i64>,
#[arg(long)]
modified_before_ms: Option<i64>,
},
/// Show the indexed file and content counts.
Stats {
@@ -48,15 +66,17 @@ enum Command {
fn main() -> Result<()> {
match Cli::parse().command {
Command::Index { db, root } => {
Command::Index { db, root, exclude } => {
let mut options = IndexOptions::default();
options.excluded_directory_names.extend(exclude);
let mut index = SearchIndex::open(&db)
.with_context(|| format!("Cannot open index {}", db.display()))?;
let report = index
.index_folder(&root)
.index_folder_with_options(&root, &options)
.with_context(|| format!("Cannot index folder {}", root.display()))?;
println!(
"Discovered: {} Indexed: {} Removed: {} Skipped: {}",
report.discovered, report.indexed, report.removed, report.skipped
"Discovered: {} Indexed: {} Unchanged: {} Removed: {} Skipped: {} Extraction failures: {}",
report.discovered, report.indexed, report.unchanged, report.removed, report.skipped, report.failures
);
}
Command::Search {
@@ -64,6 +84,12 @@ fn main() -> Result<()> {
query,
limit,
json,
folder,
extension,
min_size,
max_size,
modified_after_ms,
modified_before_ms,
} => {
if query.trim().is_empty() {
bail!("Search text cannot be empty");
@@ -72,7 +98,18 @@ fn main() -> Result<()> {
let index = SearchIndex::open(&db)
.with_context(|| format!("Cannot open index {}", db.display()))?;
let results = index
.search(&query, usize::from(limit))
.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")?;
if json {
println!("{}", serde_json::to_string_pretty(&results)?);
@@ -81,6 +118,12 @@ fn main() -> Result<()> {
} else {
for result in results {
println!("{} [{}]", result.path, result.match_kind);
if let Some(source) = result.source {
println!(" {} {}", source.kind, source.number);
}
if let Some(error) = result.extraction_error {
println!(" Text extraction unavailable: {error}");
}
if !result.snippet.is_empty() {
println!(" {}", result.snippet.replace(['\r', '\n'], " "));
}
+39
View File
@@ -88,3 +88,42 @@ fn errors_are_nonzero_and_missing_database_is_not_created() {
assert!(!bad_root.status.success());
assert!(String::from_utf8_lossy(&bad_root.stderr).contains("missing-folder"));
}
#[test]
fn exclusions_and_filters_change_real_results() {
let fixture = Fixture::new();
fs::create_dir_all(fixture.0.join("documents/private")).unwrap();
fs::write(fixture.0.join("documents/report.txt"), "quarterly invoice").unwrap();
fs::write(fixture.0.join("documents/report.md"), "quarterly invoice").unwrap();
fs::write(fixture.0.join("documents/private/secret.txt"), "invoice").unwrap();
let indexed = fixture.run(&[
"index",
"--db",
"index.db",
"--root",
"documents",
"--exclude",
"private",
]);
assert!(indexed.status.success());
let searched = fixture.run(&[
"search",
"--db",
"index.db",
"invoice",
"--extension",
"md",
"--json",
]);
assert!(
searched.status.success(),
"{}",
String::from_utf8_lossy(&searched.stderr)
);
let results: serde_json::Value = serde_json::from_slice(&searched.stdout).unwrap();
assert_eq!(results.as_array().unwrap().len(), 1);
assert_eq!(results[0]["filename"], "report.md");
let stats = fixture.run(&["stats", "--db", "index.db", "--json"]);
let stats: serde_json::Value = serde_json::from_slice(&stats.stdout).unwrap();
assert_eq!(stats["files"], 2);
}
+4
View File
@@ -9,6 +9,10 @@ anyhow.workspace = true
serde.workspace = true
rusqlite = { version = "0.32", features = ["bundled"] }
walkdir = "2"
lopdf = { version = "0.45", default-features = false }
zip = { version = "8", default-features = false, features = ["deflate"] }
quick-xml = "0.42"
[dev-dependencies]
tempfile = "3"
serde_json.workspace = true
+49 -3
View File
@@ -4,10 +4,14 @@ 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 2. Current-schema opens do
`SearchIndex::open` creates or migrates schema version 4. 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 supported-file read error rolls back the entire reconciliation.
a traversal or plain-text file read error rolls back the entire reconciliation.
PDF/DOCX parser or limit failures retain a filename record with a persisted
`SearchResult.extraction_error`, clear stale content, and increment
`IndexReport.failures`. This counts attempts in this reconciliation, so an
unchanged failed record retains its reason without incrementing the count again.
`index_folder` applies `IndexOptions::default()`: directories named `.git`,
`node_modules`, and `target` are pruned. `index_folder_with_options` lets callers
@@ -23,6 +27,46 @@ or size. Unsupported files remain searchable by filename. Content reads are
limited to 2 MiB. Reconciliation tracks seen files in SQLite rather than retaining
every path in process memory.
PDF text extraction uses `lopdf` with strict parsing and bounded load/page/font
stream decompression. DOCX uses `zip` and namespace-aware `quick-xml` to read the
main document body, including table paragraphs. The extractor does not execute
macros or follow external document relationships. It supports UTF-8 DOCX XML;
UTF-16 document parts, nested paragraphs, and XML document types are explicit
failures. Headers, footers, comments, drawings, and tracked-deletion semantics
are not yet covered. Encrypted PDFs are explicitly unsupported. Scanned PDFs
without a text layer remain filename searchable; OCR is not implemented.
| Limit | Value |
| --- | --- |
| PDF/DOCX file input, before parsing | 32 MiB |
| PDF load stream / page content / font decoding | 4 MiB per stream or page budget |
| PDF page count, after parsing | 1,000 |
| DOCX main XML, including decompressed ZIP output | 8 MiB |
| DOCX ZIP entries / paragraphs | 10,000 each |
| DOCX XML depth | 128 |
| Extracted text per document | 2 MiB |
Input reads and DOCX XML decompression are capped while reading. PDF's loader and
text extractor use the library's bounded decoding APIs. Page/entry counts and
total PDF extracted text are checked after the parser produces them. These are
not global process memory or CPU-time limits: multiple PDF objects/fonts can each
consume their budget, and parsing may still be expensive. Process isolation and
timeouts remain required before handling arbitrary untrusted documents in a
production release.
Actual page/paragraph text is stored separately in an FTS passage index. Search
results use the best passage that matches all literal query terms and expose
`source: { kind, number }`, with one-based PDF page, DOCX paragraph, or plain-text
line-range start. Paragraph numbers include empty paragraphs. A filename-only
match or a match spanning several passages can have no source; the engine does
not invent a location. Snippets with a source are drawn from that passage.
`passage_batch(after_id, limit)` exposes at most 256 `IndexedPassage` records for
real embedding runtimes. IDs represent current rows and can be reused after
reconciliation; downstream embedding caches must validate content and model
fingerprints. Consumers must tokenize long pages/paragraphs into overlapping
windows rather than silently truncate them to a model token limit.
`search_with_options` supports canonical folder boundaries, extension, inclusive
size bounds, and inclusive Unix millisecond modification-time bounds. Filters
apply before limiting results. `search` retains the unfiltered API. Query text is
@@ -46,4 +90,6 @@ escaped as literal FTS terms, not interpreted as operators. Results are capped a
- Extraction runs in the caller and a large reconciliation holds one transaction.
Bounded worker isolation, persistent jobs, and crash-resumable scans remain
production milestones.
- PDF, DOCX, OCR, media extraction, and real model inference remain unimplemented.
- OCR, media extraction, and real model inference belong to subsequent modules.
- Passage selection currently adds a bounded query per result; query batching and
passage ranking benchmarks remain necessary for large result lists.
+267
View File
@@ -0,0 +1,267 @@
//! Bounded inputs and outputs; parser allocations and CPU are not process isolated.
use crate::SourceLocation;
use anyhow::{bail, Context, Result};
use quick_xml::{events::Event, name::ResolveResult, NsReader, XmlVersion};
use std::{
fs::File,
io::{Cursor, Read},
path::Path,
};
const MAX_TEXT: usize = 2 * 1024 * 1024;
const MAX_DOCUMENT: usize = 32 * 1024 * 1024;
const MAX_XML: usize = 8 * 1024 * 1024;
const MAX_PDF_STREAM: usize = 4 * 1024 * 1024;
const MAX_PASSAGES: usize = 10_000;
pub(crate) struct Passage {
pub text: String,
pub source: SourceLocation,
}
pub(crate) fn extract(path: &Path) -> Result<Option<Vec<Passage>>> {
let extension = path
.extension()
.and_then(|value| value.to_str())
.unwrap_or("")
.to_ascii_lowercase();
if extension == "pdf" || extension == "docx" {
let bytes = read_bounded(path, MAX_DOCUMENT)?;
return if extension == "pdf" {
pdf(&bytes).map(Some)
} else {
docx(&bytes).map(Some)
};
}
if !matches!(
extension.as_str(),
"txt"
| "md"
| "rs"
| "py"
| "js"
| "ts"
| "tsx"
| "jsx"
| "json"
| "toml"
| "yaml"
| "yml"
| "csv"
| "html"
| "css"
| "xml"
| "log"
| "sh"
) {
return Ok(None);
}
let file = File::open(path)?;
if file.metadata()?.len() > MAX_TEXT as u64 {
return Ok(None);
}
let mut bytes = Vec::new();
file.take(MAX_TEXT as u64 + 1).read_to_end(&mut bytes)?;
if bytes.len() > MAX_TEXT || bytes.contains(&0) {
return Ok(None);
}
let Ok(text) = String::from_utf8(bytes) else {
return Ok(None);
};
let lines = text.split_inclusive('\n').collect::<Vec<_>>();
Ok(Some(
lines
.chunks(64)
.enumerate()
.map(|(index, lines)| Passage {
text: lines.concat(),
source: SourceLocation {
kind: "line".into(),
number: (index * 64 + 1) as u32,
},
})
.collect(),
))
}
fn read_bounded(path: &Path, maximum: usize) -> Result<Vec<u8>> {
let file = File::open(path)?;
if file.metadata()?.len() > maximum as u64 {
bail!("document input exceeds {maximum} bytes");
}
let mut bytes = Vec::new();
file.take(maximum as u64 + 1).read_to_end(&mut bytes)?;
if bytes.len() > maximum {
bail!("document input exceeds {maximum} bytes");
}
Ok(bytes)
}
fn pdf(bytes: &[u8]) -> Result<Vec<Passage>> {
let document = lopdf::Document::load_mem_with_options(
bytes,
lopdf::LoadOptions {
strict: true,
max_decompressed_size: Some(MAX_PDF_STREAM),
..Default::default()
},
)
.context("invalid, encrypted, or over-limit PDF")?;
if document.was_encrypted() || document.is_encrypted() {
bail!("encrypted PDF is unsupported");
}
let pages = document.get_pages();
if pages.is_empty() {
bail!("PDF has no pages");
}
if pages.len() > 1000 {
bail!("PDF exceeds 1000 pages");
}
let mut passages = Vec::new();
let mut total = 0;
for number in pages.keys() {
let text = document
.extract_text_with_limit(&[*number], MAX_PDF_STREAM)
.with_context(|| format!("PDF text extraction failed on page {number}"))?;
total += text.len();
if total > MAX_TEXT {
bail!("PDF text exceeds {MAX_TEXT} bytes");
}
if !text.trim().is_empty() {
passages.push(Passage {
text,
source: SourceLocation {
kind: "page".into(),
number: *number,
},
});
}
}
Ok(passages)
}
fn docx(bytes: &[u8]) -> Result<Vec<Passage>> {
let mut archive =
zip::ZipArchive::new(Cursor::new(bytes)).context("invalid DOCX ZIP container")?;
if archive.len() > 10_000 {
bail!("DOCX exceeds 10000 archive entries");
}
let mut xml = String::new();
let part = archive
.by_name("word/document.xml")
.context("DOCX lacks word/document.xml")?;
if part.size() > MAX_XML as u64 {
bail!("DOCX document XML exceeds {MAX_XML} bytes");
}
part.take(MAX_XML as u64 + 1)
.read_to_string(&mut xml)
.context("DOCX XML must be valid UTF-8")?;
if xml.len() > MAX_XML {
bail!("DOCX document XML exceeds {MAX_XML} bytes");
}
let mut reader = NsReader::from_str(&xml);
let mut passages = Vec::new();
let mut paragraph = String::new();
let mut number = 0;
let mut in_paragraph = false;
let mut in_text = false;
let mut total = 0;
let mut depth = 0usize;
let mut saw_document = false;
loop {
let (namespace, event) = reader.read_resolved_event().context("invalid DOCX XML")?;
let word_namespace = matches!(namespace, ResolveResult::Bound(value) if matches!(value.as_ref(),
"http://schemas.openxmlformats.org/wordprocessingml/2006/main" | "http://purl.oclc.org/ooxml/wordprocessingml/main"));
match &event {
Event::Start(start) => {
if depth == 0 {
if saw_document || !word_namespace || start.local_name().as_ref() != "document"
{
bail!("invalid DOCX document root");
}
saw_document = true;
}
depth += 1;
if depth > 128 {
bail!("DOCX XML exceeds 128 nesting levels");
}
}
Event::End(_) => depth = depth.checked_sub(1).context("invalid DOCX XML depth")?,
Event::Empty(_) if depth == 0 => bail!("invalid DOCX document root"),
_ => {}
}
match event {
Event::Start(event) if word_namespace => match event.local_name().as_ref() {
"p" => {
if in_paragraph {
bail!("nested DOCX paragraphs are unsupported");
}
number += 1;
if number > MAX_PASSAGES {
bail!("DOCX exceeds 10000 paragraphs");
}
in_paragraph = true;
paragraph.clear();
}
"t" => in_text = in_paragraph,
_ => {}
},
Event::End(event) if word_namespace => match event.local_name().as_ref() {
"t" => in_text = false,
"p" => {
in_paragraph = false;
in_text = false;
total += paragraph.len();
if total > MAX_TEXT {
bail!("DOCX text exceeds {MAX_TEXT} bytes");
}
if !paragraph.trim().is_empty() {
passages.push(Passage {
text: std::mem::take(&mut paragraph),
source: SourceLocation {
kind: "paragraph".into(),
number: number as u32,
},
});
}
}
_ => {}
},
Event::Empty(event) if word_namespace => match event.local_name().as_ref() {
"tab" if in_paragraph => paragraph.push('\t'),
"br" | "cr" if in_paragraph => paragraph.push('\n'),
"p" => {
number += 1;
if number > MAX_PASSAGES {
bail!("DOCX exceeds 10000 paragraphs");
}
}
_ => {}
},
Event::Text(event) if in_text => {
paragraph.push_str(&event.xml_content(XmlVersion::Implicit1_0))
}
Event::CData(event) if in_text => {
paragraph.push_str(&event.xml_content(XmlVersion::Implicit1_0))
}
Event::GeneralRef(event) if in_text => {
let reference = format!("&{};", event.as_ref());
paragraph.push_str(
&quick_xml::escape::unescape(&reference).context("unsupported DOCX entity")?,
);
}
Event::DocType(_) => bail!("DOCX XML document types are unsupported"),
Event::Eof => {
if in_paragraph || depth != 0 || !saw_document {
bail!("incomplete DOCX document");
}
break;
}
_ => {}
}
if total + paragraph.len() > MAX_TEXT {
bail!("DOCX text exceeds {MAX_TEXT} bytes");
}
}
Ok(passages)
}
+120 -56
View File
@@ -1,21 +1,20 @@
//! Offline lexical search engine shared by the CLI and desktop shell.
//!
//! Reconciliation is atomic. Incomplete directory walks never delete records.
//! Text extraction is deliberately bounded; binary document extractors belong in
//! isolated workers in the next milestone, not in this indexing transaction.
//! Text and document extraction have byte limits, but still run in this process.
//! Worker isolation remains a production milestone.
mod extraction;
use anyhow::{bail, Context, Result};
use rusqlite::{params, Connection, OptionalExtension};
use serde::Serialize;
use std::{
fs,
io::Read,
path::{Path, PathBuf},
time::UNIX_EPOCH,
};
use walkdir::WalkDir;
const SCHEMA_VERSION: i64 = 2;
const MAX_TEXT_BYTES: u64 = 2 * 1024 * 1024;
const SCHEMA_VERSION: i64 = 4;
const MAX_RESULTS: usize = 1000;
#[derive(Debug, Default, Serialize)]
@@ -25,6 +24,7 @@ pub struct IndexReport {
pub removed: usize,
pub skipped: usize,
pub unchanged: usize,
pub failures: usize,
}
/// Names of directories to prune anywhere below an explicitly chosen root.
@@ -60,6 +60,27 @@ pub struct SearchResult {
pub snippet: String,
pub score: f64,
pub match_kind: String,
pub source: Option<SourceLocation>,
pub extraction_error: Option<String>,
}
/// A verified matching passage location; numbers are one based.
#[derive(Debug, Clone, Serialize)]
pub struct SourceLocation {
pub kind: String,
pub number: u32,
}
/// A bounded corpus page for downstream real embedding runtimes. IDs identify
/// current index rows, not immutable content; consumers must validate content
/// fingerprints across reconciliations before reusing persisted embeddings.
#[derive(Debug, Serialize)]
pub struct IndexedPassage {
pub id: i64,
pub path: String,
pub filename: String,
pub text: String,
pub source: SourceLocation,
}
#[derive(Debug, Serialize)]
@@ -103,7 +124,7 @@ impl SearchIndex {
id INTEGER PRIMARY KEY, path TEXT NOT NULL UNIQUE,
root TEXT NOT NULL, filename TEXT NOT NULL, content TEXT NOT NULL,
size INTEGER NOT NULL DEFAULT -1, mtime INTEGER NOT NULL DEFAULT -1,
extension TEXT NOT NULL DEFAULT '');
extension TEXT NOT NULL DEFAULT '', extraction_error TEXT);
CREATE INDEX IF NOT EXISTS files_root ON files(root);
CREATE VIRTUAL TABLE IF NOT EXISTS search USING fts5(filename, path, content,
content='files', content_rowid='id', tokenize='unicode61');
@@ -128,9 +149,21 @@ impl SearchIndex {
ALTER TABLE files ADD COLUMN extension TEXT NOT NULL DEFAULT '';
PRAGMA user_version=2;",
)?;
} else {
initialization.execute_batch("PRAGMA user_version=2;")?;
}
if version > 0 && version < 4 {
initialization.execute("ALTER TABLE files ADD COLUMN extraction_error TEXT", [])?;
}
initialization.execute_batch(
"CREATE VIRTUAL TABLE IF NOT EXISTS passages USING fts5(
text, file_path UNINDEXED, kind UNINDEXED, number UNINDEXED, tokenize='unicode61');
CREATE TRIGGER IF NOT EXISTS files_passages_delete AFTER DELETE ON files BEGIN
DELETE FROM passages WHERE file_path=old.path; END;",
)?;
if version != 4 {
// Re-extract records from older schemas to populate genuine locations.
initialization.execute("UPDATE files SET mtime=-1", [])?;
}
initialization.execute_batch("PRAGMA user_version=4;")?;
initialization.commit()?;
}
Ok(Self {
@@ -218,17 +251,46 @@ impl SearchIndex {
report.unchanged += 1;
continue;
}
let content = extract_text(entry.path()).with_context(|| format!("read {path}"))?;
if content.is_none() {
let (extracted, extraction_error) = match extraction::extract(entry.path()) {
Ok(extracted) => (extracted, None),
Err(error) if matches!(extension.as_str(), "pdf" | "docx") => {
report.failures += 1;
let reason: String = format!("{error:#}").chars().take(2048).collect();
(None, Some(reason))
}
Err(error) => {
return Err(error)
.with_context(|| format!("read {path}; reconciliation rolled back"))
}
};
if extracted.is_none() {
report.skipped += 1;
}
let extracted = extracted.unwrap_or_default();
let content = extracted
.iter()
.map(|passage| passage.text.as_str())
.collect::<Vec<_>>()
.join("\n");
transaction.execute(
"INSERT INTO files(path,root,filename,content,size,mtime,extension) VALUES (?1,?2,?3,?4,?5,?6,?7)
"INSERT INTO files(path,root,filename,content,size,mtime,extension,extraction_error) VALUES (?1,?2,?3,?4,?5,?6,?7,?8)
ON CONFLICT(path) DO UPDATE SET root=excluded.root, filename=excluded.filename,
content=excluded.content, size=excluded.size, mtime=excluded.mtime,
extension=excluded.extension",
params![path, root_key, filename, content.unwrap_or_default(), size, mtime, extension],
extension=excluded.extension, extraction_error=excluded.extraction_error",
params![path, root_key, filename, content, size, mtime, extension, extraction_error],
)?;
transaction.execute("DELETE FROM passages WHERE file_path=?1", [&path])?;
for passage in extracted {
transaction.execute(
"INSERT INTO passages(text,file_path,kind,number) VALUES (?1,?2,?3,?4)",
params![
passage.text,
path,
passage.source.kind,
passage.source.number
],
)?;
}
report.indexed += 1;
}
report.removed = transaction.execute(
@@ -315,7 +377,7 @@ impl SearchIndex {
FROM search WHERE search MATCH ?1)
SELECT f.path, f.filename, COALESCE(h.excerpt, ''), COALESCE(h.rank, 0),
CASE WHEN lower(f.filename)=lower(?2) THEN 0
WHEN instr(lower(f.filename), lower(?2)) > 0 THEN 1 ELSE 2 END AS tier
WHEN instr(lower(f.filename), lower(?2)) > 0 THEN 1 ELSE 2 END AS tier, f.extraction_error
FROM files f LEFT JOIN hits h ON h.id=f.id
WHERE (h.id IS NOT NULL OR instr(lower(f.filename), lower(?2)) > 0)
AND (?4 IS NULL OR substr(f.path,1,length(?4))=?4)
@@ -324,7 +386,7 @@ impl SearchIndex {
AND (?8 IS NULL OR f.mtime/1000000>=?8) AND (?9 IS NULL OR f.mtime/1000000<=?9)
ORDER BY tier, h.rank, f.path LIMIT ?3",
)?;
let results = statement
let mut results = statement
.query_map(
params![
fts,
@@ -350,13 +412,56 @@ impl SearchIndex {
_ => "content",
}
.to_owned(),
source: None,
extraction_error: row.get(5)?,
})
},
)?
.collect::<rusqlite::Result<Vec<_>>>()?;
let mut passages = self.connection.prepare("SELECT kind,number,snippet(passages,0,'','',' … ',24)
FROM passages WHERE passages MATCH ?1 AND file_path=?2 ORDER BY bm25(passages),rowid LIMIT 1")?;
for result in &mut results {
let passage: Option<(String, u32, String)> = passages
.query_row(params![fts, result.path], |row| {
Ok((row.get(0)?, row.get(1)?, row.get(2)?))
})
.optional()?;
if let Some((kind, number, snippet)) = passage {
result.source = Some(SourceLocation { kind, number });
result.snippet = snippet;
}
}
Ok(results)
}
pub fn contains_path(&self, path: &str) -> Result<bool> {
Ok(self.connection.query_row(
"SELECT EXISTS(SELECT 1 FROM files WHERE path=?1)",
[path],
|row| row.get(0),
)?)
}
pub fn passage_batch(&self, after_id: i64, limit: usize) -> Result<Vec<IndexedPassage>> {
let mut statement = self.connection.prepare("SELECT p.rowid,f.path,f.filename,p.text,p.kind,p.number
FROM passages p JOIN files f ON f.path=p.file_path WHERE p.rowid>?1 ORDER BY p.rowid LIMIT ?2")?;
let passages = statement
.query_map(params![after_id, limit.min(256) as i64], |row| {
Ok(IndexedPassage {
id: row.get(0)?,
path: row.get(1)?,
filename: row.get(2)?,
text: row.get(3)?,
source: SourceLocation {
kind: row.get(4)?,
number: row.get(5)?,
},
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(passages)
}
pub fn stats(&self) -> Result<IndexStats> {
Ok(self.connection.query_row(
"SELECT count(*), count(NULLIF(content, '')) FROM files",
@@ -371,47 +476,6 @@ impl SearchIndex {
}
}
fn extract_text(path: &Path) -> Result<Option<String>> {
let extension = path
.extension()
.and_then(|value| value.to_str())
.unwrap_or("")
.to_ascii_lowercase();
if !matches!(
extension.as_str(),
"txt"
| "md"
| "rs"
| "py"
| "js"
| "ts"
| "tsx"
| "jsx"
| "json"
| "toml"
| "yaml"
| "yml"
| "csv"
| "html"
| "css"
| "xml"
| "log"
| "sh"
) {
return Ok(None);
}
let file = fs::File::open(path)?;
if file.metadata()?.len() > MAX_TEXT_BYTES {
return Ok(None);
}
let mut bytes = Vec::new();
file.take(MAX_TEXT_BYTES + 1).read_to_end(&mut bytes)?;
if bytes.len() as u64 > MAX_TEXT_BYTES || bytes.contains(&0) {
return Ok(None);
}
Ok(String::from_utf8(bytes).ok())
}
#[cfg(test)]
mod tests {
use super::*;
+196
View File
@@ -0,0 +1,196 @@
use anyhow::Result;
use lopdf::{
content::{Content, Operation},
dictionary, Document, Object, Stream,
};
use searchi_core::SearchIndex;
use std::{fs, io::Write, path::Path};
fn pdf(path: &Path) -> Result<()> {
pdf_with_text(
path,
&[
"First introductory page",
"Supplier cancellation clause on second page",
],
)
}
fn pdf_with_text(path: &Path, texts: &[&str]) -> Result<()> {
let mut document = Document::with_version("1.5");
let pages_id = document.new_object_id();
let font = document.add_object(dictionary! {"Type"=>"Font", "Subtype"=>"Type1", "BaseFont"=>"Helvetica", "Encoding"=>"WinAnsiEncoding"});
let resources = document.add_object(dictionary! {"Font"=>dictionary! {"F1"=>font}});
let mut pages = Vec::new();
for text in texts {
let content = Content {
operations: vec![
Operation::new("BT", vec![]),
Operation::new("Tf", vec![Object::Name(b"F1".to_vec()), 12.into()]),
Operation::new("Td", vec![50.into(), 700.into()]),
Operation::new("Tj", vec![Object::string_literal(*text)]),
Operation::new("ET", vec![]),
],
};
let stream = document.add_object(Stream::new(dictionary! {}, content.encode()?));
let page = document
.add_object(dictionary! {"Type"=>"Page", "Parent"=>pages_id, "Contents"=>stream});
pages.push(Object::Reference(page));
}
document.objects.insert(pages_id, Object::Dictionary(dictionary! {"Type"=>"Pages", "Kids"=>pages,
"Count"=>texts.len() as i64, "Resources"=>resources, "MediaBox"=>vec![0.into(),0.into(),595.into(),842.into()]}));
let catalog = document.add_object(dictionary! {"Type"=>"Catalog", "Pages"=>pages_id});
document.trailer.set("Root", catalog);
document.compress();
document.save(path)?;
Ok(())
}
#[test]
fn compressed_document_limits_are_enforced() -> Result<()> {
let root = tempfile::tempdir()?;
let text = "a".repeat(4 * 1024 * 1024 + 100);
pdf_with_text(&root.path().join("bomb.pdf"), &[&text])?;
let xml = format!(
r#"<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"><w:body><w:p><w:r><w:t>{}</w:t></w:r></w:p></w:body></w:document>"#,
"b".repeat(2 * 1024 * 1024 + 1)
);
docx(&root.path().join("over-limit.docx"), &xml)?;
let mut index = SearchIndex::open(":memory:")?;
assert_eq!(index.index_folder(root.path())?.failures, 2);
assert_eq!(index.stats()?.content_files, 0);
assert!(index.search("bomb.pdf", 10)?[0].extraction_error.is_some());
assert!(index.search("over-limit.docx", 10)?[0]
.extraction_error
.as_ref()
.unwrap()
.contains("exceeds"));
Ok(())
}
#[test]
fn encrypted_pdf_is_explicitly_unsupported() -> Result<()> {
let root = tempfile::tempdir()?;
let path = root.path().join("protected.pdf");
pdf(&path)?;
let mut document = Document::load(&path)?;
document.trailer.set(
"ID",
vec![
Object::string_literal("fixture-id"),
Object::string_literal("fixture-id"),
],
);
let state = lopdf::EncryptionState::try_from(lopdf::EncryptionVersion::V1 {
document: &document,
owner_password: "owner",
user_password: "reader",
permissions: lopdf::Permissions::all(),
})?;
document.encrypt(&state)?;
document.save(&path)?;
let mut index = SearchIndex::open(":memory:")?;
assert_eq!(index.index_folder(root.path())?.failures, 1);
assert!(index.search("protected.pdf", 10)?[0]
.extraction_error
.as_ref()
.unwrap()
.contains("encrypted"));
assert!(index.search("cancellation", 10)?.is_empty());
Ok(())
}
fn docx(path: &Path, xml: &str) -> Result<()> {
let mut writer = zip::ZipWriter::new(fs::File::create(path)?);
let options = zip::write::SimpleFileOptions::default()
.compression_method(zip::CompressionMethod::Deflated);
for (name, content) in [
(
"[Content_Types].xml",
r#"<?xml version="1.0"?><Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types"><Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/><Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/></Types>"#,
),
(
"_rels/.rels",
r#"<?xml version="1.0"?><Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships"><Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="word/document.xml"/></Relationships>"#,
),
("word/document.xml", xml),
] {
writer.start_file(name, options)?;
writer.write_all(content.as_bytes())?;
}
writer.finish()?;
Ok(())
}
#[test]
fn real_pdf_and_docx_return_matching_locations_and_persist() -> Result<()> {
let root = tempfile::tempdir()?;
pdf(&root.path().join("agreement.pdf"))?;
docx(
&root.path().join("notes.docx"),
r#"<?xml version="1.0" encoding="UTF-8"?><w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"><w:body><w:p><w:r><w:t>Introduction</w:t></w:r></w:p><w:p/><w:p><w:r><w:t>Research &amp; </w:t></w:r><w:r><w:t>development budget</w:t></w:r></w:p></w:body></w:document>"#,
)?;
let db = root.path().join("index.db");
let mut index = SearchIndex::open(&db)?;
assert_eq!(index.index_folder(root.path())?.indexed, 2);
let result = index.search("cancellation", 10)?;
assert_eq!(result[0].source.as_ref().unwrap().kind, "page");
assert_eq!(result[0].source.as_ref().unwrap().number, 2);
assert!(result[0].snippet.contains("cancellation"));
let result = index.search("development", 10)?;
assert_eq!(result[0].source.as_ref().unwrap().kind, "paragraph");
assert_eq!(result[0].source.as_ref().unwrap().number, 3);
assert!(result[0].snippet.contains("Research & development"));
let batch = index.passage_batch(0, 2)?;
assert_eq!(batch.len(), 2);
let remainder = index.passage_batch(batch[1].id, 256)?;
assert_eq!(remainder.len(), 2);
assert!(remainder
.iter()
.any(|passage| passage.source.kind == "paragraph" && passage.source.number == 3));
assert!(index.contains_path(&result[0].path)?);
assert!(!index.contains_path("unknown")?);
drop(index);
let index = SearchIndex::open(&db)?;
assert_eq!(
index.search("budget", 10)?[0]
.source
.as_ref()
.unwrap()
.number,
3
);
Ok(())
}
#[test]
fn corrupt_document_reports_failure_and_preserves_filename_search() -> Result<()> {
let root = tempfile::tempdir()?;
fs::write(root.path().join("notes.txt"), "persistent evidence")?;
let mut index = SearchIndex::open(":memory:")?;
index.index_folder(root.path())?;
fs::write(root.path().join("broken.pdf"), b"not a PDF")?;
assert_eq!(index.index_folder(root.path())?.failures, 1);
let result = index.search("broken.pdf", 10)?;
assert!(result[0].extraction_error.as_ref().unwrap().contains("PDF"));
assert_eq!(index.search("evidence", 10)?.len(), 1);
fs::remove_file(root.path().join("broken.pdf"))?;
docx(
&root.path().join("broken.docx"),
"<w:document><w:p><w:t>broken</w:p>",
)?;
assert_eq!(index.index_folder(root.path())?.failures, 1);
assert!(index.search("broken.docx", 10)?[0]
.extraction_error
.is_some());
pdf(&root.path().join("agreement.pdf"))?;
index.index_folder(root.path())?;
assert_eq!(index.search("cancellation", 10)?.len(), 1);
fs::write(root.path().join("agreement.pdf"), "corrupted replacement")?;
assert_eq!(index.index_folder(root.path())?.failures, 1);
assert!(index.search("cancellation", 10)?.is_empty());
assert!(index.search("agreement.pdf", 10)?[0]
.extraction_error
.is_some());
Ok(())
}
+15
View File
@@ -0,0 +1,15 @@
[package]
name = "searchi-semantic"
version.workspace = true
edition.workspace = true
license.workspace = true
description = "Bounded offline CPU ONNX embedding adapter for Searchi"
[dependencies]
anyhow.workspace = true
ort = { version = "=2.0.0-rc.13", default-features = false, features = ["std", "load-dynamic", "api-27"] }
sha2 = "0.10"
tokenizers = { version = "=0.22.2", default-features = false, features = ["fancy-regex"] }
[dev-dependencies]
serde_json.workspace = true
+290
View File
@@ -0,0 +1,290 @@
//! 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;
/// 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());
}
}
@@ -0,0 +1,71 @@
//! Opt-in real-artifact test. Default tests never acquire model/runtime assets.
use searchi_semantic::{CpuEmbedder, DIMENSIONS};
use std::path::PathBuf;
#[test]
#[ignore = "Requires explicit SEARCHI_MODEL_DIR and SEARCHI_ORT_LIBRARY local assets"]
fn real_cpu_model_retrieves_english_and_german_fixture() {
let model =
PathBuf::from(std::env::var_os("SEARCHI_MODEL_DIR").expect("Set SEARCHI_MODEL_DIR"));
let runtime =
PathBuf::from(std::env::var_os("SEARCHI_ORT_LIBRARY").expect("Set SEARCHI_ORT_LIBRARY"));
let fixture: serde_json::Value =
serde_json::from_str(include_str!("../../../tools/semantic-proof/fixture.json")).unwrap();
let documents = fixture["documents"].as_array().unwrap();
let texts: Vec<_> = documents
.iter()
.map(|doc| doc["text"].as_str().unwrap())
.collect();
let mut embedder = CpuEmbedder::load(&model, &runtime, 1).unwrap();
let vectors = embedder.encode(&texts).unwrap();
assert!(vectors.iter().all(|vector| vector.len() == DIMENSIONS));
for vector in &vectors {
let norm = vector.iter().map(|v| v * v).sum::<f32>().sqrt();
assert!((norm - 1.0).abs() < 1e-5);
}
let mut correct = 0;
for query in fixture["queries"].as_array().unwrap() {
let vector = embedder
.encode(&[query["text"].as_str().unwrap()])
.unwrap()
.remove(0);
let mut scores: Vec<_> = vectors
.iter()
.enumerate()
.map(|(i, candidate)| {
(
i,
candidate
.iter()
.zip(&vector)
.map(|(a, b)| a * b)
.sum::<f32>(),
)
})
.collect();
scores.sort_by(|a, b| b.1.total_cmp(&a.1));
let rank = scores
.iter()
.position(|(i, _)| documents[*i]["id"] == query["expected"])
.unwrap()
+ 1;
println!(
"{}: expected rank {rank}; top result {}",
query["language"], documents[scores[0].0]["id"]
);
assert!(
rank <= 3,
"Real model lost expected passage for {}",
query["text"]
);
correct += usize::from(rank == 1);
}
// This is the measured proof baseline, not the production quality gate.
// A known product-release query currently returns the expected result second.
assert!(
correct >= 5,
"Semantic proof regressed below measured 5/6 top-1 baseline"
);
assert!(embedder.encode(&[" "]).is_err());
assert!(embedder.encode(&["x"; 9]).is_err());
}