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
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:
1 parent
f77112db97
commit
a9b6592330
46 files changed
+10794
-77
No files matched your search
@@ -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`).
|
||||
@@ -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'], " "));
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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
|
||||
@@ -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.
|
||||
@@ -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
@@ -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::*;
|
||||
|
||||
@@ -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 & </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(())
|
||||
}
|
||||
@@ -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
|
||||
@@ -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());
|
||||
}
|
||||
Reference in new issue
Block a user