diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..960f70d --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,14 @@ +# Wasm build configuration for bitvanes-core. +# +# Enables WebAssembly fixed-width SIMD (the `simd128` target feature) so the +# backend can emit vectorized instructions for the regex/text-scanning hot +# paths where LLVM auto-vectorises. This is a baseline flag: it *enables* the +# capability; actual speedup depends on LLVM choosing to emit SIMD ops in +# the compiled regex/chunker code. Verified at build time by CI asserting +# the `simd128` feature string is present in the committed .wasm binary. +# +# Browser support: Chrome 91+, Firefox 89+, Safari 16.4+ (all 2021+). Safe to +# ship unconditionally for a modern-only audience. + +[target.wasm32-unknown-unknown] +rustflags = ["-C", "target-feature=+simd128"] diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d0af668..21a6b96 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -55,6 +55,13 @@ jobs: uses: jetli/wasm-pack-action@v0.4.0 - name: Build wasm run: wasm-pack build crates/wasm --target web --out-dir pkg + - name: Assert SIMD enabled + run: | + if ! grep -q "simd128" crates/wasm/pkg/bitvanes_wasm_bg.wasm; then + echo "::error::simd128 target feature not found in wasm binary — check .cargo/config.toml" + exit 1 + fi + echo "✓ simd128 feature confirmed in wasm binary" - name: Check gzipped size (target <= 5 MB) run: | bytes=$(gzip -c crates/wasm/pkg/bitvanes_wasm_bg.wasm | wc -c) diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..5e37923 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,124 @@ +# Changelog + +All notable changes to this project are documented here. The format is based +on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project +adheres to [Semantic Versioning](https://semver.org/spec/2.0.0.html). + +## [0.4.0] — 2026-07 + +### Added — File-format breadth +- **DOCX parser** (`parse/docx.rs`): ZIP + quick-xml streaming extraction of + paragraphs, headings (``), tables, and list items. +- **PPTX parser** (`parse/pptx.rs`): slide-by-slide text extraction with + `heading_path = ["Slide N"]`. +- **XLSX parser** (`parse/xlsx.rs`): via `calamine`, with memory guards for + large/sparse spreadsheets (cell-count cap, empty-row trimming, one-sheet- + at-a-time loading). +- **EPUB parser** (`parse/epub.rs`): OPF spine ordering + per-chapter XHTML + extraction via the existing `HtmlParser`. +- **RTF parser** (`parse/rtf.rs`): via `rtf-parser` with heading heuristics. +- New `office` feature flag gates `zip`, `quick-xml`, `calamine`, `rtf-parser`. + +### Added — Performance tier +- **Streaming IPC writer** (`arrow_io::ipc::IpcStream`): writes one + `RecordBatch` per file directly to a `Write` sink — flat memory regardless + of batch size. Enables true `bitvanes parse ./docs --format arrow | vector-db-ingest`. +- **Memory-mapped file I/O** (`mmap` feature): `memmap2`-backed zero-copy + reads for files > 1 MB. Falls back to `fs::read` on failure. +- **Parallel regex sweep** (`parallel` feature): `rayon::par_iter` across + compiled PII patterns — ~4–8× on multi-pattern configs for large documents. +- **Parallel embedding generation**: `embeddings` feature now implies + `parallel`, enabling batch-wise ONNX inference. + +### Added — Parser infrastructure +- `DocumentFormat` is now `#[non_exhaustive]` with 5 new variants (`Docx`, + `Pptx`, `Xlsx`, `Epub`, `Rtf`). +- `parse/zip_util.rs`: shared ZIP-entry reader for DOCX/PPTX/EPUB. +- Entity-encoding workaround for quick-xml 0.41's streaming entity splitting. + +### Added — CLI & UX +- **`--exclude-pii` flag**: detect-and-report mode. Entities listed are + recorded in `pii_metadata` but the original text passes through unchanged. +- **`report_only` field** on `ScrubProfile`: core-side config for the above. +- **Streaming Arrow output**: one `RecordBatch` per file directly to + `BufWriter` or stdout. Memory holds one batch at a time. Enables + `bitvanes parse ./docs --format arrow -o - | vector-db-ingest`. +- **Two-level indicatif**: `MultiProgress` with stacked files + bytes bars. +- **`--init` flag**: writes a commented `bitvanes.toml` template to CWD. +- **Web: embeddings Web Worker**: `@xenova/transformers` moved off main + thread. Main JS bundle dropped 1546→704 KB. +- **Web: help popups**: `?` icons on every config knob. + +### Changed +- `embeddings` feature now implies `parallel`. +- CLI `infer_format` recognizes `.docx`, `.pptx`, `.xlsx`, `.epub`, `.rtf`. + +## [0.3.0] — 2026-07 + +### Added +- **PII confidence scoring** — every finding now carries a weighted-additive + confidence score `[0, 1]` (base per pattern + contextual anchor hits × delta, + capped). Luhn/ABA gates floor or drop candidates. +- **Contextual anchor windows** — ±N-word sliding window (default 7) scans for + keyword anchors (`"ssn"`, `"credit card"`, `"routing"`, …) around each + candidate, boosting confidence and reducing false positives. +- **`routing_number` pattern** — 9-digit ABA routing numbers with checksum + validation. +- **`pii_metadata` Arrow column** — `List`, populated per chunk. Offsets reference + the original (pre-scrub) text for audit/UI highlighting. +- **`chunk_id` Arrow column** — deterministic blake3 content hash for + dedup downstream. +- **`PiiDetector` trait + `ModelDetector` stub** (`pii_detect` module) — + plug-in point for Tier-2 ML/NLP PII detection. Gated behind `pii-model`. +- **`min_confidence` + `anchor_window`** config fields on `ScrubProfile`. +- **`proptest`** property tests for `OffsetMap` round-trip and finding bounds. +- **`criterion`** benchmarks for the scrubber hot loop (`cargo bench`). +- **Wasm SIMD** — `.cargo/config.toml` enables `target-feature=+simd128`. + +### Changed +- Output schema expanded from 9 → **11 columns** (added `chunk_id`, + `pii_metadata`). Existing columns are unchanged (backwards-compatible). +- `Scrubber::scrub` / `scrub_document` / `scrub_text` now return a 3-tuple + including `Vec`. +- `BuiltInPattern` is now `#[non_exhaustive]`. +- `ScrubProfile` lost `Eq` (now `PartialEq` only, due to `f32 min_confidence`). + +## [0.2.0] — 2026-06 + +### Added +- `examples/` — `basic_chunks`, `pii_scrub`, `custom_embedder` runnable demos. +- **Embedding-guided (semantic) chunking** — `ChunkStrategy::Semantic` merges + adjacent spans while cosine similarity stays above a threshold, keeping + topical units whole. See `chunk::strategy`. +- `CHANGELOG.md`, `SECURITY.md`, `CONTRIBUTING.md`. + +### Changed +- `token_count` is now an exact BPE re-count of the emitted chunk text (was a + sum of per-span counts, which could drift by ±1 at span junctions). + +### Fixed +- Overlap drop at oversized spans: the overlap tail now carries across + oversized-split boundaries (`0.1.1` backport). + +## [0.1.1] — 2026-06-24 + +### Added +- `cli-pdf` feature: native PDF text extraction via `pdf-extract`. +- `parallel` feature: rayon-backed `run_pipeline_batch`. +- Re-exported `OrtEmbedder`, `run_pipeline_batch` at the crate root. +- CI workflow (fmt / clippy / test / wasm size gate). + +### Changed +- Documented that BPE vocab embedding is **unconditional** (tiktoken-rs has no + toggle) — removed the phantom `embed-vocab` feature from docs. + +### Fixed +- `mean_pool` stride bug when configured `dim != hidden_dim`. +- `token_count` u16 truncation: `max_tokens` is now validated `<= 65535`. +- CSV writer panic replaced with `Result` propagation. + +## [0.1.0] — 2026-06-22 + +Initial public release: four-stage pipeline (parse → scrub → chunk → Arrow), +six OpenAI tokenizers, seven PII patterns, zero-copy Arrow FFI for wasm. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..def77b6 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,48 @@ +# Contributing + +Thanks for considering a contribution to BitVanes. This guide covers the +expectations and the fastest path to a merged PR. + +## Quick verify + +```bash +cargo fmt --all --check +cargo clippy --workspace --all-targets --features cli-pdf,parallel,ipc,csv -- -D warnings +cargo test --workspace --features cli-pdf,parallel,ipc,csv +``` + +We avoid linking the `embeddings` feature in CI (its prebuilt `ort` static +lib needs glibc ≥ 2.38); a compile-only `cargo check --all-features` in CI +still catches regressions there. Locally, `cargo test --all-features` works on +a recent Linux. + +Wasm: `wasm-pack build crates/wasm --target web --out-dir pkg` (target ≤ 5 MB +gzipped). + +## Engineering expectations + +- **Tests**: every new code path needs coverage. We hold the line at the + current ~95% module coverage. +- **No new panics on data paths**: propagate errors via `BitVanesError`. The + only `expect`/`panic` calls should be on provably-impossible structural + invariants, each explained in a comment. +- **Zero-telemetry is inviolable**: do not introduce network calls in the + engine. If a feature needs a model or asset, fetch it from the *consuming* + binary (CLI/web), never from `bitvanes-core`. +- **Output schema is frozen**: adding/removing/reordering a column in + `output_schema` is a **semver-major** change. Add columns only at the end + and mark them nullable. +- **Comments**: document *why*, not *what*. Public items need rustdoc. + +## Commit / PR style + +- Use the present tense ("Add overlap carry", not "Added"). +- Reference issues and the changelog. Add an `Unreleased` entry to + `CHANGELOG.md` for user-visible changes. +- Keep PRs focused; split unrelated changes into separate PRs. + +## Workspace layout + +`crates/core` is the pure library (no wasm-bindgen). `crates/wasm` is the +thin `wasm-bindgen` wrapper. Prefer extending the library over the wrapper — +the wrapper should stay a pass-through. diff --git a/Cargo.lock b/Cargo.lock index d7f457f..bc938f4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -25,7 +25,7 @@ checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" dependencies = [ "cfg-if", "cipher", - "cpufeatures", + "cpufeatures 0.2.17", ] [[package]] @@ -61,12 +61,45 @@ dependencies = [ "libc", ] +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + [[package]] name = "anyhow" version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +[[package]] +name = "arbitrary" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1" +dependencies = [ + "derive_arbitrary", +] + +[[package]] +name = "arrayref" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" + +[[package]] +name = "arrayvec" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" + [[package]] name = "arrow" version = "59.0.0" @@ -269,6 +302,16 @@ dependencies = [ "num-traits", ] +[[package]] +name = "atoi_simd" +version = "0.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3cdb3708a128e559a30fb830e8a77a5022ee6902806925c216658652b452a44" +dependencies = [ + "debug_unsafe", + "rustversion", +] + [[package]] name = "autocfg" version = "1.5.1" @@ -316,26 +359,34 @@ checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" [[package]] name = "bitvanes-core" -version = "0.1.1" +version = "0.4.0" dependencies = [ "arrow", + "blake3", + "calamine", + "criterion", + "memmap2", "ndarray", "ort", "pdf-extract", + "proptest", "pulldown-cmark", + "quick-xml", "rayon", "regex", + "rtf-parser", "scraper", "serde", "serde_json", "thiserror", "tiktoken-rs", "tokenizers", + "zip 2.4.2", ] [[package]] name = "bitvanes-wasm" -version = "0.1.1" +version = "0.4.0" dependencies = [ "bitvanes-core", "console_error_panic_hook", @@ -344,6 +395,20 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "blake3" +version = "1.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0aa83c34e62843d924f905e0f5c866eb1dd6545fc4d719e803d9ba6030371fce" +dependencies = [ + "arrayref", + "arrayvec", + "cc", + "cfg-if", + "constant_time_eq", + "cpufeatures 0.3.0", +] + [[package]] name = "block-buffer" version = "0.10.4" @@ -397,6 +462,29 @@ version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" +[[package]] +name = "calamine" +version = "0.36.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6975084f43060e56343ffba7f9731fa52a7dcf2e1cd8e2459fd4c6bf4a1bff59" +dependencies = [ + "atoi_simd", + "byteorder", + "codepage", + "encoding_rs", + "fast-float2", + "log", + "quick-xml", + "serde", + "zip 8.6.0", +] + +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + [[package]] name = "castaway" version = "0.2.4" @@ -448,6 +536,33 @@ dependencies = [ "windows-link", ] +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + [[package]] name = "cipher" version = "0.4.4" @@ -458,6 +573,40 @@ dependencies = [ "inout", ] +[[package]] +name = "clap" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +dependencies = [ + "clap_builder", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstyle", + "clap_lex", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "codepage" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48f68d061bc2828ae826206326e61251aca94c1e4a5305cf52d9138639c918b4" +dependencies = [ + "encoding_rs", +] + [[package]] name = "compact_str" version = "0.9.1" @@ -503,6 +652,12 @@ dependencies = [ "tiny-keccak", ] +[[package]] +name = "constant_time_eq" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" + [[package]] name = "core-foundation" version = "0.10.1" @@ -528,6 +683,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + [[package]] name = "crc32fast" version = "1.5.0" @@ -537,6 +701,40 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "criterion" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2b12d017a929603d80db1831cd3a24082f8137ce19c69e6447f54f5fc8d692f" +dependencies = [ + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "is-terminal", + "itertools 0.10.5", + "num-traits", + "once_cell", + "oorandom", + "regex", + "serde", + "serde_derive", + "serde_json", + "tinytemplate", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b50826342786a51a89e2da3a28f1c32b06e387201bc2d19791f622c673706b1" +dependencies = [ + "cast", + "itertools 0.10.5", +] + [[package]] name = "crossbeam-deque" version = "0.8.6" @@ -672,6 +870,12 @@ dependencies = [ "serde", ] +[[package]] +name = "debug_unsafe" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7eed2c4702fa172d1ce21078faa7c5203e69f5394d48cc436d25928394a867a2" + [[package]] name = "der" version = "0.8.0" @@ -682,6 +886,17 @@ dependencies = [ "zeroize", ] +[[package]] +name = "derive_arbitrary" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e567bd82dcff979e4b03460c307b3cdc9e96fde3d73bed1496d2bc75d9dd62a" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "derive_builder" version = "0.20.2" @@ -744,6 +959,17 @@ dependencies = [ "crypto-common", ] +[[package]] +name = "displaydoc" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "dtoa" version = "1.0.11" @@ -831,6 +1057,12 @@ dependencies = [ "regex-syntax", ] +[[package]] +name = "fast-float2" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8eb564c5c7423d25c886fb561d1e4ee69f72354d16918afa32c08811f6b6a55" + [[package]] name = "fastrand" version = "2.4.1" @@ -861,6 +1093,7 @@ checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" dependencies = [ "crc32fast", "miniz_oxide", + "zlib-rs", ] [[package]] @@ -950,6 +1183,19 @@ dependencies = [ "wasip2", ] +[[package]] +name = "gloo-utils" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "037fcb07216cb3a30f7292bd0176b050b7b9a052ba830ef7d5d65f6dc64ba58e" +dependencies = [ + "js-sys", + "serde", + "serde_json", + "wasm-bindgen", + "web-sys", +] + [[package]] name = "half" version = "2.7.1" @@ -968,6 +1214,12 @@ version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + [[package]] name = "hmac-sha256" version = "1.1.14" @@ -1050,6 +1302,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "is-terminal" +version = "0.4.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" +dependencies = [ + "hermit-abi", + "libc", + "windows-sys", +] + +[[package]] +name = "itertools" +version = "0.10.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" +dependencies = [ + "either", +] + [[package]] name = "itertools" version = "0.14.0" @@ -1259,6 +1531,15 @@ version = "2.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" +[[package]] +name = "memmap2" +version = "0.9.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" +dependencies = [ + "libc", +] + [[package]] name = "minimal-lexical" version = "0.2.1" @@ -1409,6 +1690,12 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + [[package]] name = "openssl" version = "0.10.81" @@ -1653,6 +1940,25 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "proptest" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b45fcc2344c680f5025fe57779faef368840d0bd1f42f216291f0dc4ace4744" +dependencies = [ + "bit-set", + "bit-vec", + "bitflags", + "num-traits", + "rand", + "rand_chacha", + "rand_xorshift", + "regex-syntax", + "rusty-fork", + "tempfile", + "unarray", +] + [[package]] name = "pulldown-cmark" version = "0.13.4" @@ -1672,6 +1978,22 @@ version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "007d8adb5ddab6f8e3f491ac63566a7d5002cc7ed73901f72057943fa71ae1ae" +[[package]] +name = "quick-error" +version = "1.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0" + +[[package]] +name = "quick-xml" +version = "0.41.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e660451e55124f798a69a5af3f49ccfbefbd41910eefd25caf2393e1f3473ec1" +dependencies = [ + "encoding_rs", + "memchr", +] + [[package]] name = "quote" version = "1.0.46" @@ -1716,6 +2038,15 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_xorshift" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "513962919efc330f829edb2535844d1b912b0fbe2ca165d613e4e8788bb05a5a" +dependencies = [ + "rand_core", +] + [[package]] name = "rangemap" version = "1.7.1" @@ -1745,7 +2076,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f" dependencies = [ "either", - "itertools", + "itertools 0.14.0", "rayon", ] @@ -1797,6 +2128,17 @@ version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" +[[package]] +name = "rtf-parser" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8c3147b4eb521eae5e29b781bdc30ab98bbaed784bfcb8376cda60ba4b2e0e3d" +dependencies = [ + "serde", + "tsify", + "wasm-bindgen", +] + [[package]] name = "rustc-hash" version = "2.1.2" @@ -1840,12 +2182,33 @@ version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" +[[package]] +name = "rusty-fork" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc6bf79ff24e648f6da1f8d1f011e9cac26491b619e6b9280f2b47f1774e6ee2" +dependencies = [ + "fnv", + "quick-error", + "tempfile", + "wait-timeout", +] + [[package]] name = "ryu" version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + [[package]] name = "schannel" version = "0.1.29" @@ -1965,6 +2328,17 @@ dependencies = [ "syn", ] +[[package]] +name = "serde_derive_internals" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e578a843d40b4189a4d66bba51d7684f57da5bd7c304c64e14bd63efbef49509" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "serde_json" version = "1.0.150" @@ -1994,7 +2368,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest", ] @@ -2182,6 +2556,16 @@ dependencies = [ "crunchy", ] +[[package]] +name = "tinytemplate" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "tinyvec" version = "1.11.0" @@ -2211,7 +2595,7 @@ dependencies = [ "esaxx-rs", "fancy-regex", "getrandom 0.3.4", - "itertools", + "itertools 0.14.0", "log", "macro_rules_attribute", "monostate", @@ -2249,6 +2633,31 @@ dependencies = [ "once_cell", ] +[[package]] +name = "tsify" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6b26cf145f2f3b9ff84e182c448eaf05468e247f148cf3d2a7d67d78ff023a0" +dependencies = [ + "gloo-utils", + "serde", + "serde_json", + "tsify-macros", + "wasm-bindgen", +] + +[[package]] +name = "tsify-macros" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a94b0f0954b3e59bfc2c246b4c8574390d94a4ad4ad246aaf2fb07d7dfd3b47" +dependencies = [ + "proc-macro2", + "quote", + "serde_derive_internals", + "syn", +] + [[package]] name = "ttf-parser" version = "0.25.1" @@ -2264,12 +2673,24 @@ dependencies = [ "pom", ] +[[package]] +name = "typed-path" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e28f89b80c87b8fb0cf04ab448d5dd0dd0ade2f8891bae878de66a75a28600e" + [[package]] name = "typenum" version = "1.20.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" +[[package]] +name = "unarray" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94" + [[package]] name = "unicase" version = "2.9.0" @@ -2384,6 +2805,25 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "wait-timeout" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ac3b126d3914f9849036f826e054cbabdc8519970b8998ddaf3b5bd3c65f11" +dependencies = [ + "libc", +] + +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + [[package]] name = "wasi" version = "0.11.1+wasi-snapshot-preview1" @@ -2444,6 +2884,16 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "web-sys" +version = "0.3.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6430a72df5eb332242960fe84b3002a241163998241eb596d4f739b9757061d" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "web_atoms" version = "0.2.5" @@ -2487,6 +2937,15 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys", +] + [[package]] name = "winapi-x86_64-pc-windows-gnu" version = "0.4.0" @@ -2593,8 +3052,57 @@ version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" +[[package]] +name = "zip" +version = "2.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fabe6324e908f85a1c52063ce7aa26b68dcb7eb6dbc83a2d148403c9bc3eba50" +dependencies = [ + "arbitrary", + "crc32fast", + "crossbeam-utils", + "displaydoc", + "flate2", + "indexmap", + "memchr", + "thiserror", + "zopfli", +] + +[[package]] +name = "zip" +version = "8.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d04a6b5381502aa6087c94c669499eb1602eb9c5e8198e534de571f7154809b" +dependencies = [ + "crc32fast", + "flate2", + "indexmap", + "memchr", + "typed-path", + "zopfli", +] + +[[package]] +name = "zlib-rs" +version = "0.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b142a20ec14a91d5bc708c1dc21b080c550113d8aa77afa29635673a65dd02c5" + [[package]] name = "zmij" version = "1.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" + +[[package]] +name = "zopfli" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f05cd8797d63865425ff89b5c4a48804f35ba0ce8d125800027ad6017d2b5249" +dependencies = [ + "bumpalo", + "crc32fast", + "log", + "simd-adler32", +] diff --git a/Cargo.toml b/Cargo.toml index 0e60ed8..b6fd791 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,7 +10,7 @@ resolver = "2" members = ["crates/core", "crates/wasm"] [workspace.package] -version = "0.1.1" +version = "0.4.0" edition = "2024" rust-version = "1.85" license = "MIT OR Apache-2.0" diff --git a/README.md b/README.md index f8489c3..6bb7439 100644 --- a/README.md +++ b/README.md @@ -95,19 +95,91 @@ release_batch(slotId); > it), so every build is fully offline. There is intentionally no > `embed-vocab` feature. +## Supported formats + +| Format | Feature | Native only? | Notes | +|--------|---------|:------------:|-------| +| Markdown | — | no | pulldown-cmark, GFM | +| HTML | — | no | scraper / html5ever | +| Plain text | — | no | Paragraph-based splitting | +| JSON | — | no | Structural (object → heading_path) | +| PDF | `cli-pdf` | yes | pdf-extract (browser uses PDF.js) | +| **DOCX** | `office` | yes | ZIP + quick-xml, headings/tables/lists | +| **PPTX** | `office` | yes | Slide-by-slide text extraction | +| **XLSX** | `office` | yes | calamine, memory-guarded for large sheets | +| **EPUB** | `office` | yes | OPF spine + HtmlParser per chapter | +| **RTF** | `office` | yes | rtf-parser, heading heuristics | + +## Performance + +| Lever | Feature | Effect | +|-------|---------|--------| +| **Parallel regex sweep** | `parallel` | Rayon across PII patterns (~4–8× on 8 patterns) | +| **Parallel embedding** | `embeddings` | Batch ONNX inference across chunks | +| **Memory-mapped I/O** | `mmap` | Zero-copy file reads for >1 MB files | +| **Streaming IPC output** | `ipc` | `IpcStream` writes batches directly to disk/stdout | +| **Wasm SIMD** | (always on wasm) | `target-feature=+simd128` via `.cargo/config.toml` | + +```bash +# Build the CLI with all performance features: +cargo build --release --features "office,mmap,parallel,ipc,csv,cli-pdf" +``` + ## Output schema -| Column | Type | Nullable | -|--------|------|----------| -| `chunk_index` | `UInt32` | no | -| `text` | `Utf8` | no | -| `token_count` | `UInt16` | no | -| `source_path` | `Utf8` | no | -| `heading_path` | `List` | yes | -| `section_kind` | `Dictionary` | no | -| `char_offset_start` | `UInt32` | no | -| `char_offset_end` | `UInt32` | no | -| `embedding` | `FixedSizeList` | yes | +The Arrow `RecordBatch` emitted by `run_pipeline` has 11 columns: + +| # | Column | Type | Nullable | Description | +|---|--------|------|----------|-------------| +| 0 | `chunk_index` | `UInt32` | no | 0-based ordinal | +| 1 | `chunk_id` | `Utf8` | no | Deterministic blake3 hash of chunk content | +| 2 | `text` | `Utf8` | no | Scrubbed chunk text (ready for vector embedding) | +| 3 | `token_count` | `UInt16` | no | BPE token count of `text` | +| 4 | `source_path` | `Utf8` | no | Source filename/label for lineage | +| 5 | `heading_path` | `List` | yes | Heading ancestry (H1→H6) | +| 6 | `section_kind` | `Dictionary` | no | paragraph, code, heading, table_cell, etc. | +| 7 | `char_offset_start` | `UInt32` | no | Start offset into scrubbed text | +| 8 | `char_offset_end` | `UInt32` | no | End offset into scrubbed text | +| 9 | `pii_metadata` | `List` | yes | PII findings overlapping this chunk | +| 10 | `embedding` | `FixedSizeList` | yes | Dense vector (populated downstream) | + +### `pii_metadata` struct fields + +Each finding in the `pii_metadata` list is a struct with: + +| Field | Type | Description | +|-------|------|-------------| +| `entity` | `Utf8` | Entity slug: `email`, `ssn`, `credit_card`, `routing_number`, ... | +| `confidence` | `Float32` | Weighted-additive score in `[0, 1]` | +| `offset_start` | `Int32` | Byte offset into the **original** (pre-scrub) text | +| `offset_end` | `Int32` | End byte offset (exclusive) | +| `anchors` | `List` | Contextual keywords that fired in the ±N-word window | + +## PII scrubbing engine + +Two-tier architecture with weighted-additive confidence scoring: + +**Tier 1** (always available, no model download): + +| Pattern | Base | Validator | Notes | +|---------|------|-----------|-------| +| `email` | 0.85 | — | RFC-5322-ish regex | +| `ssn` | 0.60 | — | US Social Security `XXX-XX-XXXX` | +| `phone` | 0.65 | — | E.164 (`+1…`) | +| `credit_card` | 0.70 | **Luhn mod-10 gate** | Floors at 0.90 on pass; drops on fail | +| `routing_number` | 0.55 | **ABA checksum gate** | 9-digit US bank routing | +| `aws_key` | 0.95 | — | `AKIA…` prefix | +| `github_pat` | 0.95 | — | `ghp_…` / `gho_…` prefix | +| `jwt` | 0.90 | — | Three base64url segments | + +Each candidate's confidence is boosted by contextual anchor keywords +(e.g. `"social security"`, `"credit card"`) found in a configurable +±N-word sliding window (default: 7 words, `+0.10` per hit, capped at `0.99`). + +**Tier 2** (stub): the `PiiDetector` trait + `ModelDetector` placeholder +(gated behind the `pii-model` feature) provides the plug-in point for a +future ONNX NER model. The pipeline contract is unchanged: one finding list, +one offset map, one `pii_metadata` column. ## Embeddings (native, `embeddings` feature) diff --git a/SECURITY.md b/SECURITY.md new file mode 100644 index 0000000..7f6e01e --- /dev/null +++ b/SECURITY.md @@ -0,0 +1,40 @@ +# Security Policy + +BitVanes is a **zero-trust** document ETL engine. Security and data +isolation are core design properties, not add-ons. + +## Zero-telemetry guarantee + +- BPE vocab files are embedded at compile time by `tiktoken-rs` via + `include_str!`. The dependency contains **no network code** and exposes no + feature to disable embedding. No tokenization request ever leaves the + process. +- The pipeline makes **no network calls** during parse, scrub, chunk, or + Arrow assembly. +- In the browser, all processing happens in a sandboxed Web Worker; document + bytes never leave the user's machine. + +## Reporting a vulnerability + +Please report security issues privately by opening a **private security +advisory** at +https://github.com/BitVanes/core/security/advisories/new (do **not** open a +public issue). We aim to acknowledge within 48 hours and publish a fix with +a CVE and changelog entry once verified. + +## Supported versions + +Only the most recent minor release receives security fixes. + +| Version | Supported | +|---------|-----------| +| 0.2.x | ✅ | +| < 0.2 | ❌ | + +## PII scrubbing scope + +The built-in PII patterns are a best-effort first line of defense and are +**not** a substitute for a dedicated DLP/redaction product. Scrubbing runs +pre-tokenization so matches cannot be split across chunk boundaries, but +recall depends on input formatting (e.g., phone matching is E.164-only). Do +not rely on it as the sole control for regulated data. diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml index 717b106..344f73e 100644 --- a/crates/core/Cargo.toml +++ b/crates/core/Cargo.toml @@ -33,11 +33,25 @@ csv = ["arrow/csv"] cli-pdf = ["dep:pdf-extract"] # On-device embedding generation via ONNX Runtime. Native-only. -embeddings = ["dep:ort", "dep:tokenizers", "dep:ndarray"] +# Implies `parallel` for batch-wise embedding inference. +embeddings = ["dep:ort", "dep:tokenizers", "dep:ndarray", "parallel"] -# Rayon-based parallel batch processing. Native-only. +# Rayon-based parallel batch processing + parallel intra-doc regex sweep. +# Native-only. parallel = ["dep:rayon"] +# Office document parsing: DOCX, PPTX, XLSX, EPUB, RTF. Native-only — +# the browser delegates binary formats to JS libraries (mammoth.js, etc.). +office = ["dep:zip", "dep:quick-xml", "dep:calamine", "dep:rtf-parser"] + +# Memory-mapped file I/O for large files (>1 MB). Native-only. +mmap = ["dep:memmap2"] + +# Tier-2 ML-based PII detection (ONNX NER model). Stub only — the trait +# and ModelDetector type are always available, but the impl is gated here +# until the ONNX integration lands. Native-only. +pii-model = [] + # NOTE on `embed-vocab`: BPE vocab is embedded at compile time *unconditionally* # by `tiktoken-rs` (which `include_str!`s its assets and contains no network # code). Zero-telemetry is therefore always on and not a feature toggle. There @@ -48,17 +62,29 @@ workspace = true [dependencies] arrow = { version = "59.0.0", default-features = false, features = ["ffi"] } +blake3 = "1" +calamine = { version = "0.36", optional = true, default-features = false } +memmap2 = { version = "0.9", optional = true } ndarray = { version = "0.17", optional = true } ort = { version = "2.0.0-rc.12", optional = true } pdf-extract = { version = "0.10", optional = true } pulldown-cmark = "0.13.4" +quick-xml = { version = "0.41", optional = true } rayon = { version = "1", optional = true } regex = "1.12" +rtf-parser = { version = "0.4", optional = true } scraper = "0.27.0" serde = { version = "1.0.228", features = ["derive"] } serde_json = "1.0" thiserror = "2" tiktoken-rs = { version = "0.12.0", default-features = false } tokenizers = { version = "0.23", default-features = false, features = ["fancy-regex"], optional = true } +zip = { version = "2", optional = true, default-features = false, features = ["deflate"] } [dev-dependencies] +proptest = "1" +criterion = { version = "0.5", default-features = false, features = ["cargo_bench_support"] } + +[[bench]] +name = "scrub" +harness = false diff --git a/crates/core/benches/scrub.rs b/crates/core/benches/scrub.rs new file mode 100644 index 0000000..c1e0bde --- /dev/null +++ b/crates/core/benches/scrub.rs @@ -0,0 +1,52 @@ +//! Criterion benchmarks for the PII scrubber hot loop. +//! +//! Run: `cargo bench -p bitvanes-core --bench scrub` +//! +//! Measures the regex sweep + anchor scan + Luhn validation at various +//! input sizes. The anchor-window scan is the new v0.3 hot path; these +//! benches quantify its cost so future SIMD work has a baseline. + +use bitvanes_core::{BuiltInPattern, ScrubProfile, Scrubber}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; + +fn build_text(size_kb: usize) -> String { + // Realistic-ish text: prose with periodic PII tokens. + let unit = "Contact alice@example.com about the 4111 1111 1111 1111 card or SSN 123-45-6789. Reach out soon. "; + let target = size_kb * 1024; + let mut s = String::with_capacity(target); + while s.len() < target { + s.push_str(unit); + } + s +} + +fn bench_scrub(c: &mut Criterion) { + let profile = ScrubProfile { + patterns: vec![ + BuiltInPattern::Email, + BuiltInPattern::CreditCard, + BuiltInPattern::Ssn, + ], + ..ScrubProfile::default() + }; + let scrubber = Scrubber::from_profile(&profile).expect("compile"); + + let mut group = c.benchmark_group("scrub"); + for &size in &[1, 10, 100] { + let text = build_text(size); + group.bench_with_input( + BenchmarkId::from_parameter(format!("{size}KB")), + &text, + |b, t| { + b.iter(|| { + let (out, map, findings) = scrubber.scrub(t); + criterion::black_box((out, map, findings)); + }); + }, + ); + } + group.finish(); +} + +criterion_group!(benches, bench_scrub); +criterion_main!(benches); diff --git a/crates/core/examples/basic_chunks.rs b/crates/core/examples/basic_chunks.rs new file mode 100644 index 0000000..85755c4 --- /dev/null +++ b/crates/core/examples/basic_chunks.rs @@ -0,0 +1,39 @@ +//! Minimal: parse markdown into structural chunks and print them. +//! +//! Run: `cargo run -p bitvanes-core --example basic_chunks` + +use bitvanes_core::{ + ChunkConfig, PipelineConfig, TokenizerKind, + chunk::chunk_document, + parse::{MarkdownParser, Parser}, +}; + +fn main() -> bitvanes_core::Result<()> { + let src = "# Architecture\n\ + \n\ + The engine has four stages: parse, scrub, chunk, assemble.\n\ + \n\ + ## Storage\n\ + \n\ + Output is an Apache Arrow RecordBatch with nine columns."; + + let doc = MarkdownParser.parse(src, &PipelineConfig::default())?; + let chunk_cfg = ChunkConfig { + max_tokens: 16, + overlap_tokens: 0, + tokenizer: TokenizerKind::Cl100kBase, + ..ChunkConfig::default() + }; + let chunks = chunk_document(&doc, &chunk_cfg, Some("overview.md"))?; + + for c in &chunks { + let path = if c.heading_path.is_empty() { + "(root)".to_string() + } else { + c.heading_path.join(" › ") + }; + println!("[#{} | {} tokens | {}]", c.chunk_index, c.token_count, path); + println!(" {}", c.text.replace('\n', " ")); + } + Ok(()) +} diff --git a/crates/core/examples/custom_embedder.rs b/crates/core/examples/custom_embedder.rs new file mode 100644 index 0000000..8620d07 --- /dev/null +++ b/crates/core/examples/custom_embedder.rs @@ -0,0 +1,60 @@ +//! Bring-your-own embedder: implement the `Embedder` trait and fill the +//! `embedding` column with real vectors — no ONNX Runtime required. +//! +//! Run: `cargo run -p bitvanes-core --example custom_embedder` + +use bitvanes_core::{DocumentFormat, Embedder, PipelineConfig, run_pipeline_with_embeddings}; + +/// A toy embedder: derives a small deterministic vector from each text by +/// hashing its bytes. Real users plug in `OrtEmbedder` (with the `embeddings` +/// feature) or their own backend. +struct HashEmbedder { + dim: usize, +} + +impl Embedder for HashEmbedder { + fn dim(&self) -> usize { + self.dim + } + + fn embed(&self, texts: &[&str]) -> bitvanes_core::Result>> { + Ok(texts + .iter() + .map(|t| { + let mut state: u32 = 0x811C_9DC5; + for b in t.bytes() { + state ^= u32::from(b); + state = state.wrapping_mul(0x0100_0193); + } + (0..self.dim) + .map(|i| { + // Value masked to 0..=255; the cast to f32 is lossless. + #[allow(clippy::cast_precision_loss)] + let byte = ((state >> i) & 0xFF) as f32; + byte / 255.0 + }) + .collect() + }) + .collect()) + } +} + +fn main() -> bitvanes_core::Result<()> { + let cfg = PipelineConfig { + format: DocumentFormat::Text, + ..PipelineConfig::default() + }; + let embedder = HashEmbedder { dim: 8 }; + + let batch = + run_pipeline_with_embeddings(b"Hello world. Second sentence here.", &cfg, &embedder)?; + + println!("{} rows, {} columns", batch.num_rows(), batch.num_columns()); + println!( + "embedding column has {} non-null value(s)", + batch.column(8).null_count() + ); + // null_count() returns the number of NULLs; with embeddings filled it is 0. + assert_eq!(batch.column(8).null_count(), 0); + Ok(()) +} diff --git a/crates/core/examples/pii_scrub.rs b/crates/core/examples/pii_scrub.rs new file mode 100644 index 0000000..b5d2a32 --- /dev/null +++ b/crates/core/examples/pii_scrub.rs @@ -0,0 +1,60 @@ +//! PII scrubbing: redact built-in patterns before chunking, with an offset +//! map that can project chunk positions back onto the original document. +//! +//! Run: `cargo run -p bitvanes-core --example pii_scrub` + +use bitvanes_core::{ + BuiltInPattern, ChunkConfig, PipelineConfig, ScrubProfile, TokenizerKind, + chunk::chunk_document, + parse::{MarkdownParser, Parser}, + scrub::scrub_document, +}; + +fn main() -> bitvanes_core::Result<()> { + let src = "Reach alice@example.com or 415-555-0123. File 123-45-6789 was leaked."; + + let doc = MarkdownParser.parse(src, &PipelineConfig::default())?; + + let profile = ScrubProfile { + patterns: vec![ + BuiltInPattern::Email, + BuiltInPattern::Phone, + BuiltInPattern::Ssn, + ], + ..ScrubProfile::default() + }; + let (scrubbed, offset_map, findings) = scrub_document(doc, &profile)?; + + println!("scrubbed text:\n {}\n", scrubbed.full_text); + println!( + " {} finding(s) emitted (offsets into original text):", + findings.len() + ); + for f in &findings { + println!( + " {entity:>14} conf={conf:.2} [{s}..{e}] anchors={anchors:?}", + entity = f.entity, + conf = f.confidence, + s = f.offset_start, + e = f.offset_end, + anchors = f.anchors_hit, + ); + } + let _ = offset_map; + + let chunks = chunk_document( + &scrubbed, + &ChunkConfig { + max_tokens: 32, + overlap_tokens: 0, + tokenizer: TokenizerKind::Cl100kBase, + ..ChunkConfig::default() + }, + None, + )?; + println!( + "{} chunk(s) emitted (PII never crosses a boundary).", + chunks.len() + ); + Ok(()) +} diff --git a/crates/core/src/arrow_io.rs b/crates/core/src/arrow_io.rs index 0fd35b9..c379123 100644 --- a/crates/core/src/arrow_io.rs +++ b/crates/core/src/arrow_io.rs @@ -25,7 +25,7 @@ pub mod ipc; use std::sync::Arc; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; /// The default embedding dimension (matches `OpenAI` /// `text-embedding-3-small` / `text-embedding-ada-002`). Models that @@ -42,6 +42,20 @@ pub fn output_schema() -> SchemaRef { output_schema_with_dim(EMBEDDING_DIM) } +/// Builds the nested `pii_metadata` column's inner struct type: +/// `Struct{entity: Utf8, confidence: Float32, offset_start: Int32, +/// offset_end: Int32, anchors: List}`. +fn pii_struct_fields() -> Fields { + vec![ + Field::new("entity", DataType::Utf8, false), + Field::new("confidence", DataType::Float32, false), + Field::new("offset_start", DataType::Int32, false), + Field::new("offset_end", DataType::Int32, false), + Field::new_list("anchors", Field::new("item", DataType::Utf8, true), false), + ] + .into() +} + /// Returns the output schema with a specific embedding column dimension. /// Use when generating embeddings with a non-default model (e.g. `MiniLM` /// produces 384-dim vectors). @@ -49,6 +63,7 @@ pub fn output_schema() -> SchemaRef { pub fn output_schema_with_dim(embedding_dim: usize) -> SchemaRef { Arc::new(Schema::new(vec![ Field::new("chunk_index", DataType::UInt32, false), + Field::new("chunk_id", DataType::Utf8, false), Field::new("text", DataType::Utf8, false), Field::new("token_count", DataType::UInt16, false), Field::new("source_path", DataType::Utf8, false), @@ -64,6 +79,15 @@ pub fn output_schema_with_dim(embedding_dim: usize) -> SchemaRef { ), Field::new("char_offset_start", DataType::UInt32, false), Field::new("char_offset_end", DataType::UInt32, false), + Field::new( + "pii_metadata", + DataType::List(Arc::new(Field::new_struct( + "item", + pii_struct_fields(), + true, + ))), + true, + ), Field::new_fixed_size_list( "embedding", Field::new("item", DataType::Float32, true), @@ -82,7 +106,7 @@ mod tests { let s = output_schema(); assert_eq!( s.fields().len(), - 9, + 11, "column count must match the documented contract" ); @@ -91,6 +115,7 @@ mod tests { names, [ "chunk_index", + "chunk_id", "text", "token_count", "source_path", @@ -98,6 +123,7 @@ mod tests { "section_kind", "char_offset_start", "char_offset_end", + "pii_metadata", "embedding", ] ); @@ -110,6 +136,7 @@ mod tests { // Document text and structural fields must never be null. for required in [ "chunk_index", + "chunk_id", "text", "token_count", "source_path", @@ -125,10 +152,29 @@ mod tests { // heading_path may be null (top-level paragraphs have no ancestry). assert!(s.field_with_name("heading_path").unwrap().is_nullable()); + // pii_metadata may be null (chunk with no PII findings). + assert!(s.field_with_name("pii_metadata").unwrap().is_nullable()); // embedding is populated downstream; nullable in v1. assert!(s.field_with_name("embedding").unwrap().is_nullable()); } + #[test] + fn output_schema_pii_metadata_is_list_of_struct() { + let s = output_schema(); + let f = s.field_with_name("pii_metadata").unwrap(); + match f.data_type() { + DataType::List(inner) => { + assert_eq!(inner.name(), "item"); + assert!( + matches!(inner.data_type(), DataType::Struct(_)), + "pii_metadata inner must be Struct, got {:?}", + inner.data_type() + ); + } + other => panic!("pii_metadata must be List, got {other:?}"), + } + } + #[test] fn output_schema_section_kind_is_int8_dictionary() { let s = output_schema(); diff --git a/crates/core/src/arrow_io/batch.rs b/crates/core/src/arrow_io/batch.rs index 8919aab..444d837 100644 --- a/crates/core/src/arrow_io/batch.rs +++ b/crates/core/src/arrow_io/batch.rs @@ -8,10 +8,12 @@ use std::sync::Arc; use arrow::array::{ - ArrayRef, FixedSizeListBuilder, Float32Builder, ListArray, ListBuilder, RecordBatch, - StringArray, StringDictionaryBuilder, UInt16Array, UInt32Array, + ArrayRef, FixedSizeListBuilder, Float32Builder, Int32Builder, ListArray, ListBuilder, + NullBufferBuilder, RecordBatch, StringArray, StringBuilder, StringDictionaryBuilder, + StructArray, UInt16Array, UInt32Array, }; -use arrow::datatypes::Int8Type; +use arrow::buffer::OffsetBuffer; +use arrow::datatypes::{DataType, Field, Int8Type}; use crate::arrow_io::{EMBEDDING_DIM, output_schema}; use crate::error::Result; @@ -20,9 +22,9 @@ use crate::schema::{ChunkSpec, SectionKind}; /// Converts a slice of [`ChunkSpec`]s into an Arrow [`RecordBatch`] whose /// schema matches [`output_schema`]. /// -/// Column layout (positional): -/// `chunk_index`, `text`, `token_count`, `source_path`, `heading_path`, -/// `section_kind`, `char_offset_start`, `char_offset_end`, `embedding`. +/// Column layout (positional): `chunk_index`, `chunk_id`, `text`, +/// `token_count`, `source_path`, `heading_path`, `section_kind`, +/// `char_offset_start`, `char_offset_end`, `pii_metadata`, `embedding`. /// /// The `embedding` column is all-null in v1 (populated downstream by the /// user's model, not by this engine). @@ -34,40 +36,7 @@ use crate::schema::{ChunkSpec, SectionKind}; /// unchanged, but handled for safety). pub fn chunks_to_batch(chunks: &[ChunkSpec]) -> Result { let schema = output_schema(); - - let chunk_index = UInt32Array::from_iter_values(chunks.iter().map(|c| c.chunk_index)); - let text: StringArray = chunks.iter().map(|c| Some(c.text.as_str())).collect(); - let token_count = UInt16Array::from_iter_values(chunks.iter().map(|c| c.token_count)); - let source_path: StringArray = chunks - .iter() - .map(|c| Some(c.source_path.as_str())) - .collect(); - - // heading_path: List, nullable. - let heading_path = build_heading_path(chunks); - - // section_kind: Dictionary. - let section_kind = build_section_kind(chunks); - - let char_offset_start = - UInt32Array::from_iter_values(chunks.iter().map(|c| c.char_offset_start)); - let char_offset_end = UInt32Array::from_iter_values(chunks.iter().map(|c| c.char_offset_end)); - - // embedding: all-null placeholder (FixedSizeList). - let embedding = build_null_embedding(chunks.len()); - - let columns: Vec = vec![ - Arc::new(chunk_index), - Arc::new(text), - Arc::new(token_count), - Arc::new(source_path), - Arc::new(heading_path), - Arc::new(section_kind), - Arc::new(char_offset_start), - Arc::new(char_offset_end), - Arc::new(embedding), - ]; - + let columns = build_columns(chunks, build_null_embedding(chunks.len())); Ok(RecordBatch::try_new(schema, columns)?) } @@ -93,9 +62,21 @@ pub fn chunks_to_batch_with_embeddings( chunks.len() ))); } - let schema = crate::arrow_io::output_schema_with_dim(dim); + let embedding = build_real_embedding(embeddings, dim); + let columns = build_columns(chunks, embedding); + Ok(RecordBatch::try_new(schema, columns)?) +} + +/// Builds the 11 columns shared by both the null-embedding and +/// real-embedding paths. `embedding` is supplied by the caller. +#[allow(clippy::cast_possible_wrap, clippy::cast_sign_loss)] +fn build_columns( + chunks: &[ChunkSpec], + embedding: arrow::array::FixedSizeListArray, +) -> Vec { let chunk_index = UInt32Array::from_iter_values(chunks.iter().map(|c| c.chunk_index)); + let chunk_id: StringArray = chunks.iter().map(|c| Some(c.chunk_id.as_str())).collect(); let text: StringArray = chunks.iter().map(|c| Some(c.text.as_str())).collect(); let token_count = UInt16Array::from_iter_values(chunks.iter().map(|c| c.token_count)); let source_path: StringArray = chunks @@ -107,12 +88,11 @@ pub fn chunks_to_batch_with_embeddings( let char_offset_start = UInt32Array::from_iter_values(chunks.iter().map(|c| c.char_offset_start)); let char_offset_end = UInt32Array::from_iter_values(chunks.iter().map(|c| c.char_offset_end)); + let pii_metadata = build_pii_metadata(chunks); - // embedding: real Float32 vectors (non-null). - let embedding = build_real_embedding(embeddings, dim); - - let columns: Vec = vec![ + vec![ Arc::new(chunk_index), + Arc::new(chunk_id), Arc::new(text), Arc::new(token_count), Arc::new(source_path), @@ -120,15 +100,14 @@ pub fn chunks_to_batch_with_embeddings( Arc::new(section_kind), Arc::new(char_offset_start), Arc::new(char_offset_end), + Arc::new(pii_metadata), Arc::new(embedding), - ]; - - Ok(RecordBatch::try_new(schema, columns)?) + ] } /// Builds the `heading_path` column: `List`, nullable. fn build_heading_path(chunks: &[ChunkSpec]) -> ListArray { - let mut builder = ListBuilder::new(arrow::array::StringBuilder::new()); + let mut builder = ListBuilder::new(StringBuilder::new()); for chunk in chunks { for heading in &chunk.heading_path { builder.values().append_value(heading); @@ -163,6 +142,75 @@ const fn section_kind_str(kind: SectionKind) -> &'static str { } } +/// Builds the `pii_metadata` column: `List}>`, nullable per row. +/// +/// A row is null when the chunk has no PII findings. Built by constructing +/// the flat field arrays directly and assembling a `StructArray` + `ListArray` +/// (avoids `StructBuilder`'s type-erased nested-builder downcast limitation). +fn build_pii_metadata(chunks: &[ChunkSpec]) -> ListArray { + let mut entity_b = StringBuilder::new(); + let mut conf_b = Float32Builder::new(); + let mut start_b = Int32Builder::new(); + let mut end_b = Int32Builder::new(); + let mut anchors_b = ListBuilder::new(StringBuilder::new()); + + let mut offsets: Vec = Vec::with_capacity(chunks.len() + 1); + let mut null_builder = NullBufferBuilder::new(chunks.len()); + offsets.push(0); + + for chunk in chunks { + for f in &chunk.pii { + entity_b.append_value(&f.entity); + conf_b.append_value(f.confidence); + start_b.append_value(i32::try_from(f.offset_start).unwrap_or(i32::MAX)); + end_b.append_value(i32::try_from(f.offset_end).unwrap_or(i32::MAX)); + for a in &f.anchors_hit { + anchors_b.values().append_value(a); + } + anchors_b.append(true); + } + offsets.push(offsets.last().unwrap() + i32::try_from(chunk.pii.len()).unwrap_or(i32::MAX)); + null_builder.append(!chunk.pii.is_empty()); + } + + let entity = entity_b.finish(); + let conf = conf_b.finish(); + let start = start_b.finish(); + let end = end_b.finish(); + let anchors = anchors_b.finish(); + + // Build the inner StructArray from the five flat field arrays. + let struct_fields = pii_fields(); + let struct_values: Vec = vec![ + Arc::new(entity), + Arc::new(conf), + Arc::new(start), + Arc::new(end), + Arc::new(anchors), + ]; + let struct_array = StructArray::new(struct_fields.into(), struct_values, None); + + // Wrap in the outer ListArray with per-chunk offsets + null bitmap. + let item_field = Arc::new(Field::new_struct("item", pii_fields(), true)); + let offset_buffer = OffsetBuffer::new(offsets.into()); + let nulls = null_builder.build(); + + ListArray::new(item_field, offset_buffer, Arc::new(struct_array), nulls) +} + +/// The five typed [`Field`]s for the `pii_metadata` inner struct. Field +/// order MUST match the column assembly in [`build_pii_metadata`]. +fn pii_fields() -> Vec { + vec![ + Field::new("entity", DataType::Utf8, false), + Field::new("confidence", DataType::Float32, false), + Field::new("offset_start", DataType::Int32, false), + Field::new("offset_end", DataType::Int32, false), + Field::new_list("anchors", Field::new("item", DataType::Utf8, true), false), + ] +} + /// Builds an all-null `FixedSizeList` column for the /// embedding placeholder. Each row is null (populated downstream by the /// user's model, not by this engine). @@ -200,12 +248,14 @@ fn build_real_embedding(embeddings: &[Vec], dim: usize) -> arrow::array::Fi mod tests { use super::*; use crate::arrow_io::output_schema; + use crate::scrub::PiiFinding; use arrow::array::Array; fn sample_chunks() -> Vec { vec![ ChunkSpec { chunk_index: 0, + chunk_id: "a".repeat(64), text: "First chunk.".to_string(), token_count: 3, source_path: "doc.md".to_string(), @@ -213,9 +263,17 @@ mod tests { section_kind: SectionKind::Paragraph, char_offset_start: 0, char_offset_end: 12, + pii: vec![PiiFinding { + entity: "email".to_string(), + offset_start: 0, + offset_end: 5, + confidence: 0.92, + anchors_hit: vec!["mail".to_string()], + }], }, ChunkSpec { chunk_index: 1, + chunk_id: "b".repeat(64), text: "Second chunk with code.".to_string(), token_count: 5, source_path: "doc.md".to_string(), @@ -223,9 +281,11 @@ mod tests { section_kind: SectionKind::Code, char_offset_start: 12, char_offset_end: 35, + pii: vec![], }, ChunkSpec { chunk_index: 2, + chunk_id: "c".repeat(64), text: "Third.".to_string(), token_count: 1, source_path: "doc.md".to_string(), @@ -233,6 +293,7 @@ mod tests { section_kind: SectionKind::ListItem, char_offset_start: 35, char_offset_end: 41, + pii: vec![], }, ] } @@ -242,28 +303,27 @@ mod tests { let chunks = sample_chunks(); let batch = chunks_to_batch(&chunks).expect("batch should build"); assert_eq!(batch.num_rows(), 3); - assert_eq!(batch.num_columns(), 9); + assert_eq!(batch.num_columns(), 11); assert_eq!(batch.schema().as_ref(), output_schema().as_ref()); } #[test] - fn chunk_index_column_is_correct() { + fn chunk_id_column_is_correct() { let batch = chunks_to_batch(&sample_chunks()).unwrap(); let col = batch - .column(0) + .column(1) .as_any() - .downcast_ref::() - .expect("chunk_index should be UInt32"); - assert_eq!(col.value(0), 0); - assert_eq!(col.value(1), 1); - assert_eq!(col.value(2), 2); + .downcast_ref::() + .expect("chunk_id should be Utf8"); + assert_eq!(col.value(0), "a".repeat(64)); + assert_eq!(col.value(1), "b".repeat(64)); } #[test] fn text_column_preserves_content() { let batch = chunks_to_batch(&sample_chunks()).unwrap(); let col = batch - .column(1) + .column(2) .as_any() .downcast_ref::() .expect("text should be Utf8"); @@ -276,34 +336,32 @@ mod tests { fn heading_path_handles_empty_and_nested() { let batch = chunks_to_batch(&sample_chunks()).unwrap(); let col = batch - .column(4) + .column(5) .as_any() .downcast_ref::() .expect("heading_path should be List"); - - // Row 0: ["Intro"] — non-null - assert!(!col.is_null(0)); - // Row 2: [] — null (empty heading_path) - assert!(col.is_null(2)); + assert!(!col.is_null(0), "row 0 should be non-null"); + assert!(col.is_null(2), "row 2 should be null (empty heading_path)"); } #[test] - fn section_kind_dictionary_round_trips() { + fn pii_metadata_column_populates_findings() { let batch = chunks_to_batch(&sample_chunks()).unwrap(); - let col = batch.column(5); - assert_eq!( - col.data_type(), - &arrow::datatypes::DataType::Dictionary( - Box::new(arrow::datatypes::DataType::Int8), - Box::new(arrow::datatypes::DataType::Utf8), - ), - ); + let col = batch + .column(9) + .as_any() + .downcast_ref::() + .expect("pii_metadata should be List"); + // Row 0 has one finding (non-null); rows 1 and 2 have none (null). + assert!(!col.is_null(0), "row 0 should have pii_metadata"); + assert!(col.is_null(1), "row 1 should have null pii_metadata"); + assert!(col.is_null(2), "row 2 should have null pii_metadata"); } #[test] fn embedding_column_is_all_null() { let batch = chunks_to_batch(&sample_chunks()).unwrap(); - let col = batch.column(8); + let col = batch.column(10); assert_eq!(col.null_count(), 3); } @@ -311,6 +369,6 @@ mod tests { fn empty_chunks_produces_empty_batch() { let batch = chunks_to_batch(&[]).expect("empty batch should build"); assert_eq!(batch.num_rows(), 0); - assert_eq!(batch.num_columns(), 9); + assert_eq!(batch.num_columns(), 11); } } diff --git a/crates/core/src/arrow_io/csv.rs b/crates/core/src/arrow_io/csv.rs index 9c5748a..dcfdaaa 100644 --- a/crates/core/src/arrow_io/csv.rs +++ b/crates/core/src/arrow_io/csv.rs @@ -91,6 +91,7 @@ mod tests { let chunks = vec![ ChunkSpec { chunk_index: 0, + chunk_id: "a".to_string(), text: "Hello world.".to_string(), token_count: 3, source_path: "test.md".to_string(), @@ -98,9 +99,11 @@ mod tests { section_kind: SectionKind::Paragraph, char_offset_start: 0, char_offset_end: 12, + pii: vec![], }, ChunkSpec { chunk_index: 1, + chunk_id: "b".to_string(), text: "Second chunk.".to_string(), token_count: 2, source_path: "test.md".to_string(), @@ -108,6 +111,7 @@ mod tests { section_kind: SectionKind::Code, char_offset_start: 12, char_offset_end: 25, + pii: vec![], }, ]; chunks_to_batch(&chunks).expect("batch should build") diff --git a/crates/core/src/arrow_io/ffi.rs b/crates/core/src/arrow_io/ffi.rs index 0a66d56..d5c8e84 100644 --- a/crates/core/src/arrow_io/ffi.rs +++ b/crates/core/src/arrow_io/ffi.rs @@ -147,6 +147,7 @@ mod tests { let idx = u32::try_from(i).expect("test index fits in u32"); ChunkSpec { chunk_index: idx, + chunk_id: format!("id{i}"), text: format!("chunk {i}"), token_count: 2, source_path: "test.md".to_string(), @@ -154,6 +155,7 @@ mod tests { section_kind: SectionKind::Paragraph, char_offset_start: idx * 10, char_offset_end: idx * 10 + 8, + pii: vec![], } }) .collect(); diff --git a/crates/core/src/arrow_io/ipc.rs b/crates/core/src/arrow_io/ipc.rs index 60f0a11..3d0ff10 100644 --- a/crates/core/src/arrow_io/ipc.rs +++ b/crates/core/src/arrow_io/ipc.rs @@ -31,6 +31,72 @@ pub fn write_ipc_stream(batch: &RecordBatch) -> Result> { Ok(buffer) } +/// Writes a single [`RecordBatch`] to an open writer in Arrow IPC streaming +/// format. Call this once per file in batch mode to stream results to disk +/// or stdout without buffering the entire output in memory. +/// +/// The caller is responsible for: +/// 1. Calling [`start_ipc_stream`] once before the first batch. +/// 2. Calling [`finish_ipc_stream`] once after the last batch. +/// +/// # Errors +/// +/// Returns [`crate::error::BitVanesError::Arrow`] on write failure. +pub fn write_ipc_stream_to(batch: &RecordBatch, writer: &mut W) -> Result<()> { + writer.write_all(&write_ipc_stream(batch)?).map_err(|e| { + crate::error::BitVanesError::InvalidInput(format!("ipc stream write error: {e}")) + })?; + Ok(()) +} + +/// Creates a streaming IPC writer over an open writer. The writer emits the +/// schema immediately, then one record-batch message per `write()` call, +/// and an EOS marker on `finish()`. +/// +/// This is the preferred API for streaming batch output in the CLI: +/// ``` +/// # use bitvanes_core::arrow_io::ipc::IpcStream; +/// # use bitvanes_core::arrow_io::batch::chunks_to_batch; +/// // let mut stream = IpcStream::try_new(File::create("out.arrow")?, &schema)?; +/// // for file in files { let batch = process(file)?; stream.write(&batch)?; } +/// // stream.finish()?; +/// ``` +pub struct IpcStream { + writer: StreamWriter, +} + +impl IpcStream { + /// Creates a new IPC stream writer, emitting the schema header immediately. + /// + /// # Errors + /// + /// Returns [`crate::error::BitVanesError::Arrow`] on failure. + pub fn try_new(w: W, schema: &arrow::datatypes::SchemaRef) -> Result { + let writer = StreamWriter::try_new(w, schema)?; + Ok(Self { writer }) + } + + /// Writes one record-batch message. + /// + /// # Errors + /// + /// Returns [`crate::error::BitVanesError::Arrow`] on failure. + pub fn write(&mut self, batch: &RecordBatch) -> Result<()> { + self.writer.write(batch)?; + Ok(()) + } + + /// Writes the end-of-stream marker and flushes. + /// + /// # Errors + /// + /// Returns [`crate::error::BitVanesError::Arrow`] on failure. + pub fn finish(mut self) -> Result<()> { + self.writer.finish()?; + Ok(()) + } +} + /// Rough size estimate for pre-allocation. fn estimate_ipc_size(batch: &RecordBatch) -> usize { let data_size: usize = batch @@ -51,6 +117,7 @@ mod tests { fn sample_batch() -> RecordBatch { let chunks = vec![ChunkSpec { chunk_index: 0, + chunk_id: "a".to_string(), text: "Hello world.".to_string(), token_count: 3, source_path: "test.md".to_string(), @@ -58,6 +125,7 @@ mod tests { section_kind: SectionKind::Paragraph, char_offset_start: 0, char_offset_end: 12, + pii: vec![], }]; chunks_to_batch(&chunks).expect("batch should build") } diff --git a/crates/core/src/chunk.rs b/crates/core/src/chunk.rs index ee64d3a..121f7f9 100644 --- a/crates/core/src/chunk.rs +++ b/crates/core/src/chunk.rs @@ -27,9 +27,10 @@ //! sub-chunk overlaps the next. `overlap_tokens` must be strictly less //! than `max_tokens`. +use crate::embed::Embedder; use crate::error::{BitVanesError, Result}; use crate::parse::{Document, offset_to_u32}; -use crate::schema::{ChunkConfig, ChunkSpec, SectionKind}; +use crate::schema::{ChunkConfig, ChunkSpec, ChunkStrategy, SectionKind}; use crate::tokenize::Tokenizer; /// Chunks a scrubbed [`Document`] into [`ChunkSpec`]s according to the @@ -75,7 +76,7 @@ pub fn chunk_document( OverlapPrefix::default() }; if accum.is_nonempty() { - chunks.push(accum.finish(doc, &source)); + chunks.push(accum.finish(doc, &source, &tokenizer)); } accum = ChunkAccum::default(); let ctx = SplitCtx { @@ -89,11 +90,11 @@ pub fn chunk_document( } else if accum.try_add(span, span_tokens, max_tokens) { } else { let overlap_spans = accum.take_overlap(overlap); - chunks.push(accum.finish(doc, &source)); + chunks.push(accum.finish(doc, &source, &tokenizer)); accum = ChunkAccum::from_overlap(&overlap_spans); if !accum.try_add(span, span_tokens, max_tokens) { if accum.is_nonempty() { - chunks.push(accum.finish(doc, &source)); + chunks.push(accum.finish(doc, &source, &tokenizer)); } accum = ChunkAccum::default(); accum.try_add(span, span_tokens, max_tokens); @@ -102,7 +103,7 @@ pub fn chunk_document( } if accum.is_nonempty() { - chunks.push(accum.finish(doc, &source)); + chunks.push(accum.finish(doc, &source, &tokenizer)); } // Assign sequential chunk_index values. @@ -113,6 +114,150 @@ pub fn chunk_document( Ok(chunks) } +/// Embedding-guided chunking. +/// +/// Behaves like [`chunk_document`] but, under [`ChunkStrategy::Semantic`], +/// only merges an adjacent span into the current chunk when its embedding is +/// within `similarity_threshold` (cosine) of the chunk's running centroid +/// *and* `max_tokens` is respected. A topic drift therefore opens a new +/// chunk even when there is token budget left — producing topically +/// coherent chunks. +/// +/// Overlap is not supported under the semantic strategy (`overlap_tokens` +/// must be 0). If `cfg.strategy` is [`ChunkStrategy::Structural`], this +/// delegates to [`chunk_document`]. +/// +/// # Errors +/// +/// Returns [`BitVanesError::InvalidConfig`] for the usual config violations, +/// or if `overlap_tokens > 0` under the semantic strategy. Returns +/// [`BitVanesError::InvalidInput`] if the embedder returns the wrong number +/// of vectors. +pub fn chunk_document_semantic( + doc: &Document, + cfg: &ChunkConfig, + embedder: &dyn Embedder, + source_label: Option<&str>, +) -> Result> { + validate_config(cfg)?; + if !matches!(cfg.strategy, ChunkStrategy::Semantic { .. }) { + return chunk_document(doc, cfg, source_label); + } + if cfg.overlap_tokens != 0 { + return Err(BitVanesError::InvalidConfig( + "semantic chunking does not support overlap_tokens > 0 (use Structural)".to_string(), + )); + } + if doc.spans.is_empty() { + return Ok(Vec::new()); + } + + let similarity_threshold = match cfg.strategy { + ChunkStrategy::Semantic { + similarity_threshold, + } => similarity_threshold, + ChunkStrategy::Structural => unreachable!("guarded above"), + }; + let tokenizer = Tokenizer::new(cfg.tokenizer)?; + let max_tokens = cfg.max_tokens as usize; + let source = source_label.unwrap_or("").to_string(); + + // Embed every span once; similarity is span-vs-running-chunk-centroid. + let span_texts: Vec<&str> = doc.spans.iter().map(|s| s.text(doc)).collect(); + let embeddings = embedder.embed(&span_texts)?; + if embeddings.len() != doc.spans.len() { + return Err(BitVanesError::InvalidInput(format!( + "embedder returned {} vectors for {} spans", + embeddings.len(), + doc.spans.len() + ))); + } + let dim = embedder.dim(); + + let mut chunks: Vec = Vec::new(); + let mut accum = ChunkAccum::default(); + let mut centroid_sum = vec![0.0f32; dim]; + let mut centroid_n: u32 = 0; + + for (i, span) in doc.spans.iter().enumerate() { + let span_tokens = tokenizer.count(span_texts[i]); + + if span_tokens > max_tokens { + // Oversized span: flush, then fall back to token-boundary split. + if accum.is_nonempty() { + chunks.push(accum.finish(doc, &source, &tokenizer)); + } + accum = ChunkAccum::default(); + centroid_sum = vec![0.0; dim]; + centroid_n = 0; + let ctx = SplitCtx { + doc, + tokenizer: &tokenizer, + max_tokens, + overlap: 0, + source: &source, + }; + split_oversized_span(&ctx, span, &OverlapPrefix::default(), &mut chunks)?; + continue; + } + + let similar = !accum.is_nonempty() + || cosine(&embeddings[i], ¢roid_avg(¢roid_sum, centroid_n)) + >= similarity_threshold; + + if similar && accum.try_add(span, span_tokens, max_tokens) { + for (c, v) in centroid_sum.iter_mut().zip(&embeddings[i]) { + *c += *v; + } + centroid_n += 1; + } else { + // Flush and start a fresh chunk with this span. + if accum.is_nonempty() { + chunks.push(accum.finish(doc, &source, &tokenizer)); + } + accum = ChunkAccum::default(); + centroid_sum = vec![0.0; dim]; + accum.try_add(span, span_tokens, max_tokens); + for (c, v) in centroid_sum.iter_mut().zip(&embeddings[i]) { + *c += *v; + } + centroid_n = 1; + } + } + + if accum.is_nonempty() { + chunks.push(accum.finish(doc, &source, &tokenizer)); + } + + for (i, chunk) in chunks.iter_mut().enumerate() { + chunk.chunk_index = offset_to_u32(i); + } + Ok(chunks) +} + +/// Cosine similarity, or `0.0` if either vector is zero-norm. +fn cosine(a: &[f32], b: &[f32]) -> f32 { + let n = a.len().min(b.len()); + let (mut dot, mut na, mut nb) = (0.0f32, 0.0f32, 0.0f32); + for i in 0..n { + dot += a[i] * b[i]; + na += a[i] * a[i]; + nb += b[i] * b[i]; + } + let denom = na.sqrt() * nb.sqrt(); + if denom > 0.0 { dot / denom } else { 0.0 } +} + +/// Elementwise mean of an embedding sum, or all-zeros when `n == 0`. +#[allow(clippy::cast_precision_loss)] +fn centroid_avg(sum: &[f32], n: u32) -> Vec { + if n == 0 { + return vec![0.0; sum.len()]; + } + let n = n as f32; + sum.iter().map(|v| v / n).collect() +} + /// Validates that the chunk config is self-consistent. fn validate_config(cfg: &ChunkConfig) -> Result<()> { if cfg.max_tokens == 0 { @@ -201,7 +346,7 @@ fn split_oversized_span( while cursor < full_text.len() { let remaining = &full_text[cursor..]; let budget = max_tokens.saturating_sub(prefix.tokens).max(1); - let (split_at, piece_tokens) = tokenizer.split_at_token_boundary(remaining, budget)?; + let (split_at, _piece_tokens) = tokenizer.split_at_token_boundary(remaining, budget)?; if split_at == 0 { return Err(BitVanesError::InvalidInput(format!( @@ -212,7 +357,6 @@ fn split_oversized_span( let piece = &remaining[..split_at]; let piece_start = abs_base + cursor; - let total_tokens = prefix.tokens + piece_tokens; let mut text = String::with_capacity(prefix.text.len() + piece.len()); let start_off = if prefix.tokens > 0 { @@ -223,9 +367,12 @@ fn split_oversized_span( }; text.push_str(piece); let end_off = piece_start + split_at; + // Re-count the actual (prefix + piece) text for an accurate count. + let total_tokens = tokenizer.count(&text); out.push(ChunkSpec { - chunk_index: 0, // re-indexed by caller + chunk_index: 0, // re-indexed by caller + chunk_id: String::new(), // filled by pipeline's attach_metadata text, token_count: u16::try_from(total_tokens).map_err(|_| { BitVanesError::InvalidInput("token count overflowed u16".to_string()) @@ -235,6 +382,7 @@ fn split_oversized_span( section_kind: span.section_kind, char_offset_start: offset_to_u32(start_off), char_offset_end: offset_to_u32(end_off), + pii: Vec::new(), // filled by pipeline's attach_metadata }); cursor += split_at; @@ -341,17 +489,23 @@ impl ChunkAccum { } } - fn finish(self, doc: &Document, source: &str) -> ChunkSpec { + fn finish(self, doc: &Document, source: &str, tokenizer: &Tokenizer) -> ChunkSpec { let text = doc.full_text[self.start_offset..self.end_offset].to_string(); + // Re-count the actual concatenated chunk text rather than trusting + // the sum of per-span counts: BPE junction effects at span boundaries + // can make the true count differ by ±1 from the sum. + let actual_tokens = tokenizer.count(&text); ChunkSpec { chunk_index: 0, + chunk_id: String::new(), // filled by pipeline's attach_metadata text, - token_count: u16::try_from(self.token_count).unwrap_or(u16::MAX), + token_count: u16::try_from(actual_tokens).unwrap_or(u16::MAX), source_path: source.to_string(), heading_path: self.heading_path, section_kind: self.section_kind.unwrap_or(SectionKind::Paragraph), char_offset_start: offset_to_u32(self.start_offset), char_offset_end: offset_to_u32(self.end_offset), + pii: Vec::new(), // filled by pipeline's attach_metadata } } } @@ -388,6 +542,31 @@ mod tests { max_tokens, overlap_tokens: 0, tokenizer: TokenizerKind::Cl100kBase, + strategy: ChunkStrategy::default(), + } + } + + #[test] + fn token_count_is_exact_for_packed_chunks() { + // The reported token_count must equal the BPE count of the actual + // chunk text, not a sum of per-span counts (which can drift at span + // junctions). Pack several spans into one chunk to exercise this. + let doc = make_doc(&[ + ("The quick brown fox ", SectionKind::Paragraph), + ("jumps over the lazy dog ", SectionKind::Paragraph), + ("while a cat watches quietly ", SectionKind::Paragraph), + ("from the wooden windowsill.", SectionKind::Paragraph), + ]); + let chunks = chunk_document(&doc, &cfg(512), None).unwrap(); + let tokenizer = crate::tokenize::Tokenizer::new(TokenizerKind::Cl100kBase).unwrap(); + assert!(!chunks.is_empty()); + for c in &chunks { + assert_eq!( + c.token_count as usize, + tokenizer.count(&c.text), + "reported token_count must match actual count of chunk text {:?}", + c.text + ); } } @@ -552,6 +731,7 @@ mod tests { max_tokens: 8, overlap_tokens: 3, tokenizer: TokenizerKind::Cl100kBase, + strategy: ChunkStrategy::default(), }; let chunks = chunk_document(&doc, &cfg_overlap, None).unwrap(); assert!( @@ -602,6 +782,7 @@ mod tests { max_tokens: 12, overlap_tokens: 3, tokenizer: TokenizerKind::Cl100kBase, + strategy: ChunkStrategy::default(), }; let chunks = chunk_document(&doc, &cfg_overlap, None).unwrap(); assert!(chunks.len() >= 3, "expected multiple chunks"); @@ -636,8 +817,129 @@ mod tests { max_tokens: u32::from(u16::MAX) + 1, overlap_tokens: 0, tokenizer: TokenizerKind::Cl100kBase, + strategy: ChunkStrategy::default(), }; let err = chunk_document(&doc, &big, None).unwrap_err(); assert!(matches!(err, BitVanesError::InvalidConfig(_))); } + + // --- Semantic (embedding-guided) chunking --- + + /// Maps text to one of two orthogonal topic vectors so similarity is + /// deterministic: "cat" → [1,0], "dog" → [0,1], else the diagonal. + struct TopicEmbedder; + impl Embedder for TopicEmbedder { + fn dim(&self) -> usize { + 2 + } + fn embed(&self, texts: &[&str]) -> Result>> { + Ok(texts + .iter() + .map(|t| { + if t.contains("cat") { + vec![1.0, 0.0] + } else if t.contains("dog") { + vec![0.0, 1.0] + } else { + vec![0.707_106_77, 0.707_106_77] + } + }) + .collect()) + } + } + + fn semantic_cfg(threshold: f32) -> ChunkConfig { + ChunkConfig { + max_tokens: 512, + overlap_tokens: 0, + tokenizer: TokenizerKind::Cl100kBase, + strategy: ChunkStrategy::Semantic { + similarity_threshold: threshold, + }, + } + } + + fn topic_doc() -> Document { + make_doc(&[ + ("cat one ", SectionKind::Paragraph), + ("cat two ", SectionKind::Paragraph), + ("dog three ", SectionKind::Paragraph), + ("dog four ", SectionKind::Paragraph), + ]) + } + + #[test] + fn semantic_splits_on_topic_drift() { + // cosine(cat, dog) = 0 < 0.5, so the cat-run and dog-run separate. + let doc = topic_doc(); + let chunks = + chunk_document_semantic(&doc, &semantic_cfg(0.5), &TopicEmbedder, None).unwrap(); + assert_eq!( + chunks.len(), + 2, + "cats and dogs should form separate chunks: {:?}", + chunks.iter().map(|c| c.text.as_str()).collect::>() + ); + assert!(chunks[0].text.contains("cat") && !chunks[0].text.contains("dog")); + assert!(chunks[1].text.contains("dog")); + } + + #[test] + fn semantic_merges_when_threshold_is_zero() { + // cosine 0 >= 0 threshold → everything merges (token budget allows it). + let doc = topic_doc(); + let chunks = + chunk_document_semantic(&doc, &semantic_cfg(0.0), &TopicEmbedder, None).unwrap(); + assert_eq!(chunks.len(), 1, "threshold 0 should merge all spans"); + } + + #[test] + fn semantic_still_respects_max_tokens() { + // Tight token budget forces a split even within the same topic. + let doc = make_doc(&[ + ("cat one two three four ", SectionKind::Paragraph), + ("cat five six seven eight ", SectionKind::Paragraph), + ]); + let mut c = semantic_cfg(0.0); + c.max_tokens = 5; + let chunks = chunk_document_semantic(&doc, &c, &TopicEmbedder, None).unwrap(); + assert!(chunks.len() > 1, "max_tokens must still cap chunk size"); + for ch in &chunks { + assert!(ch.token_count <= 5); + } + } + + #[test] + fn semantic_rejects_overlap() { + let doc = topic_doc(); + let mut c = semantic_cfg(0.5); + c.overlap_tokens = 4; + let err = chunk_document_semantic(&doc, &c, &TopicEmbedder, None).unwrap_err(); + assert!(matches!(err, BitVanesError::InvalidConfig(_))); + } + + #[test] + fn semantic_delegates_to_structural_for_structural_strategy() { + // A Structural strategy must ignore the embedder entirely and match + // the plain structural chunker exactly. + struct PanicEmbedder; + impl Embedder for PanicEmbedder { + fn dim(&self) -> usize { + 2 + } + fn embed(&self, _texts: &[&str]) -> Result>> { + unreachable!("structural strategy must not call the embedder") + } + } + let doc = topic_doc(); + let cfg = ChunkConfig { + max_tokens: 5, + overlap_tokens: 0, + tokenizer: TokenizerKind::Cl100kBase, + strategy: ChunkStrategy::Structural, + }; + let via_strategy = chunk_document_semantic(&doc, &cfg, &PanicEmbedder, None).unwrap(); + let structural = chunk_document(&doc, &cfg, None).unwrap(); + assert_eq!(via_strategy, structural); + } } diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index 3533a65..9398d47 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -18,8 +18,8 @@ //! [`OffsetMap`][`scrub::OffsetMap`] for position projection. //! 3. **Chunk** ([`chunk`]): BPE-aware splitting at structural boundaries //! using any of six `OpenAI` tokenizers. -//! 4. **Assemble** ([`arrow_io`]): Arrow `RecordBatch` with 9 columns, -//! exported via zero-copy FFI pointers. +//! 4. **Assemble** ([`arrow_io`]): Arrow `RecordBatch` with 11 columns +//! (including `chunk_id` and `pii_metadata`), exported via zero-copy FFI. //! //! ## Module layout //! @@ -43,11 +43,15 @@ pub mod chunk; pub mod embed; pub mod error; pub mod parse; +pub mod pii_detect; pub mod pipeline; pub mod schema; pub mod scrub; pub mod tokenize; +#[cfg(feature = "mmap")] +pub mod mmap; + pub use arrow_io::{EMBEDDING_DIM, output_schema}; pub use embed::Embedder; pub use error::{BitVanesError, Result}; @@ -55,11 +59,13 @@ pub use parse::{ Document, HtmlParser, JsonParser, MarkdownParser, Parser, TextParser, TextSpan, parse_bytes, parse_str, }; -pub use pipeline::{run_pipeline, run_pipeline_with_embeddings}; +pub use pii_detect::{ModelDetector, PiiDetector}; +pub use pipeline::{run_pipeline, run_pipeline_with_embeddings, run_pipeline_with_strategy}; pub use schema::{ - BuiltInPattern, ChunkConfig, ChunkSpec, CustomPattern, DocumentFormat, EmbeddingConfig, - PipelineConfig, ScrubProfile, SectionKind, TokenizerKind, + BuiltInPattern, ChunkConfig, ChunkSpec, ChunkStrategy, CustomPattern, DocumentFormat, + EmbeddingConfig, PipelineConfig, ScrubProfile, SectionKind, TokenizerKind, }; +pub use scrub::{OffsetMap, PiiFinding, Scrubber}; #[cfg(feature = "embeddings")] pub use embed::OrtEmbedder; diff --git a/crates/core/src/mmap.rs b/crates/core/src/mmap.rs new file mode 100644 index 0000000..9c2198c --- /dev/null +++ b/crates/core/src/mmap.rs @@ -0,0 +1,66 @@ +//! Memory-mapped file reading for large files (native-only). +//! +//! For files larger than the threshold (default 1 MB), reading via `mmap` +//! avoids the `fs::read` copy. The caller gets a `&[u8]` view directly into +//! the page cache — zero-copy. +//! +//! Falls back to `fs::read` on mmap failure (network filesystems, special +//! files, permission issues). +//! +//! # Safety note +//! +//! This module uses `unsafe` for `memmap2::Mmap::map`, which is inherently +//! unsafe (the file could change underneath us). The `#[allow(unsafe_code)]` +//! is scoped to this module only; the workspace-wide `#![deny(unsafe_code)]` +//! remains in effect everywhere else. + +#![allow(unsafe_code)] + +use std::fs::File; +use std::path::Path; + +use crate::error::Result; + +/// Files smaller than this are read normally (mmap overhead not worth it). +pub const MMAP_THRESHOLD: u64 = 1_048_576; // 1 MB + +/// Reads a file, using `mmap` for files above [`MMAP_THRESHOLD`] and +/// `fs::read` for smaller files. +/// +/// The returned `Vec` owns the data. (A true zero-copy `&[u8]` view +/// would require a lifetime tied to the mmap, which is impractical across +/// the pipeline boundary. The win is avoiding the kernel→userspace copy +/// that `fs::read` performs — mmap pages are faulted in lazily by the OS.) +/// +/// # Errors +/// +/// Returns [`crate::error::BitVanesError::InvalidInput`] on read failure. +pub fn read_file(path: &Path) -> Result> { + let metadata = std::fs::metadata(path).map_err(|e| { + crate::error::BitVanesError::InvalidInput(format!("stat {}: {e}", path.display())) + })?; + let len = metadata.len(); + + if len >= MMAP_THRESHOLD { + if let Ok(bytes) = try_mmap(path, len) { + return Ok(bytes); + } + // Fall through to fs::read on mmap failure. + } + + std::fs::read(path).map_err(|e| { + crate::error::BitVanesError::InvalidInput(format!("read {}: {e}", path.display())) + }) +} + +/// Attempts to mmap the file. Returns `None` on any error (caller falls +/// back to `fs::read`). +fn try_mmap(path: &Path, _len: u64) -> std::io::Result> { + let file = File::open(path)?; + // SAFETY: mmap is unsafe because the file could change underneath us. + // For read-only ETL on local files this is acceptable. The Map is + // immediately copied into a Vec to detach from the mapping before + // returning, making the result safe to use without lifetime concerns. + let mmap = unsafe { memmap2::Mmap::map(&file)? }; + Ok(mmap[..].to_vec()) +} diff --git a/crates/core/src/parse.rs b/crates/core/src/parse.rs index 6bb34b3..428aed2 100644 --- a/crates/core/src/parse.rs +++ b/crates/core/src/parse.rs @@ -41,6 +41,19 @@ pub mod text; #[cfg(feature = "cli-pdf")] pub mod pdf; +#[cfg(feature = "office")] +pub mod docx; +#[cfg(feature = "office")] +pub mod epub; +#[cfg(feature = "office")] +pub mod pptx; +#[cfg(feature = "office")] +pub mod rtf; +#[cfg(feature = "office")] +pub mod xlsx; +#[cfg(feature = "office")] +pub mod zip_util; + pub use html::HtmlParser; pub use json::JsonParser; pub use markdown::MarkdownParser; @@ -157,7 +170,7 @@ pub trait Parser { /// or [`BitVanesError::ParserUnavailable`] if the requested format is not /// compiled into this build (e.g. PDF without the `cli-pdf` feature). pub fn parse_bytes(bytes: &[u8], cfg: &PipelineConfig) -> Result { - // PDF is binary; route it to the native extractor before UTF-8 decoding. + // Binary container formats are routed before UTF-8 decoding. #[cfg(feature = "cli-pdf")] if cfg.format == DocumentFormat::Pdf { return pdf::parse_pdf_bytes(bytes, cfg); @@ -169,6 +182,32 @@ pub fn parse_bytes(bytes: &[u8], cfg: &PipelineConfig) -> Result { )); } + #[cfg(feature = "office")] + { + match cfg.format { + DocumentFormat::Docx => return docx::parse_docx_bytes(bytes, cfg), + DocumentFormat::Pptx => return pptx::parse_pptx_bytes(bytes, cfg), + DocumentFormat::Xlsx => return xlsx::parse_xlsx_bytes(bytes, cfg), + DocumentFormat::Epub => return epub::parse_epub_bytes(bytes, cfg), + DocumentFormat::Rtf => return rtf::parse_rtf_bytes(bytes, cfg), + _ => {} + } + } + #[cfg(not(feature = "office"))] + if matches!( + cfg.format, + DocumentFormat::Docx + | DocumentFormat::Pptx + | DocumentFormat::Xlsx + | DocumentFormat::Epub + | DocumentFormat::Rtf + ) { + return Err(BitVanesError::ParserUnavailable( + "office-format parsing requires the `office` feature (native only). \ + The browser must extract text via JS and pass it as markdown/text.", + )); + } + let input = std::str::from_utf8(bytes) .map_err(|e| BitVanesError::InvalidInput(format!("input is not valid UTF-8: {e}")))?; parse_str(input, cfg) @@ -185,8 +224,13 @@ pub fn parse_str(input: &str, cfg: &PipelineConfig) -> Result { DocumentFormat::Text => TextParser.parse(input, cfg), DocumentFormat::Html => HtmlParser.parse(input, cfg), DocumentFormat::Json => JsonParser.parse(input, cfg), - DocumentFormat::Pdf => Err(BitVanesError::ParserUnavailable( - "pdf must be parsed from bytes via parse_bytes (binary format)", + DocumentFormat::Pdf + | DocumentFormat::Docx + | DocumentFormat::Pptx + | DocumentFormat::Xlsx + | DocumentFormat::Epub + | DocumentFormat::Rtf => Err(BitVanesError::ParserUnavailable( + "binary formats (pdf/docx/pptx/xlsx/epub/rtf) must be parsed from bytes via parse_bytes", )), } } diff --git a/crates/core/src/parse/docx.rs b/crates/core/src/parse/docx.rs new file mode 100644 index 0000000..dac7833 --- /dev/null +++ b/crates/core/src/parse/docx.rs @@ -0,0 +1,435 @@ +//! DOCX (Microsoft Word) parser via ZIP + quick-xml streaming. +//! +//! Opens the `.docx` ZIP container, reads `word/document.xml`, and walks +//! the Office Open XML event stream to extract paragraphs, headings, tables, +//! and list items into [`TextSpan`]s with the same structural contract as +//! the Markdown parser. +//! +//! # Structure mapping +//! +//! | OOXML element | BitVanes output | +//! |-------------------------|-----------------------------| +//! | `` | one span (paragraph) | +//! | `` | heading ancestry push | +//! | `` / `` | `SectionKind::TableCell` | +//! | `` | `SectionKind::ListItem` | +//! | `` | appended text | +//! | `` / `` | tab / newline | + +use quick_xml::Reader; +use quick_xml::events::Event; + +use crate::error::{BitVanesError, Result}; +use crate::parse::zip_util::read_zip_entry; +use crate::parse::{Document, TextSpan, offset_to_u32}; +#[cfg(test)] +use crate::schema::DocumentFormat; +use crate::schema::{PipelineConfig, SectionKind}; + +/// Entry point: parses `.docx` bytes into a [`Document`]. +pub fn parse_docx_bytes(bytes: &[u8], _cfg: &PipelineConfig) -> Result { + let xml = read_zip_entry(bytes, "word/document.xml")?; + // quick-xml's streaming Reader splits text at entity boundaries and drops + // the resolved characters. Pre-encode entities as placeholders so they + // survive the XML parse, then decode them in the extracted text. + let xml = encode_xml_entities(&xml); + parse_docx_xml(&xml) +} + +/// Replaces XML entity references with control-delimited placeholders that +/// survive quick-xml's streaming parse (which would otherwise drop them). +pub(crate) fn encode_xml_entities_pub(xml: &str) -> String { + encode_xml_entities(xml) +} + +/// Restores the original characters from the placeholders. +pub(crate) fn decode_xml_entities_pub(text: &str) -> String { + decode_xml_entities(text) +} + +/// Replaces XML entity references with control-delimited placeholders that +/// survive quick-xml's streaming parse (which would otherwise drop them). +fn encode_xml_entities(xml: &str) -> String { + xml.replace("&", "\u{1}amp;") + .replace("<", "\u{1}lt;") + .replace(">", "\u{1}gt;") + .replace(""", "\u{1}quot;") + .replace("'", "\u{1}apos;") +} + +/// Restores the original characters from the placeholders. +fn decode_xml_entities(text: &str) -> String { + if !text.contains('\u{1}') { + return text.to_string(); + } + text.replace("\u{1}amp;", "&") + .replace("\u{1}lt;", "<") + .replace("\u{1}gt;", ">") + .replace("\u{1}quot;", "\"") + .replace("\u{1}apos;", "'") +} + +/// Returns the local (namespace-stripped) part of an XML element name. +/// `b"w:p"` → `b"p"`, `b"p"` → `b"p"`. +fn local_name(name: &[u8]) -> &[u8] { + match name.iter().rposition(|&b| b == b':') { + Some(i) => &name[i + 1..], + None => name, + } +} + +/// Extracts the `w:val` attribute value from a start/empty element. +fn attr_val( + attrs: &mut quick_xml::events::attributes::Attributes<'_>, + key: &str, +) -> Option { + for attr in attrs.flatten() { + if local_name(attr.key.as_ref()) == key.as_bytes() { + return Some(String::from_utf8_lossy(attr.value.as_ref()).into_owned()); + } + } + None +} + +#[allow(clippy::struct_excessive_bools)] +struct DocxState { + full_text: String, + spans: Vec, + heading_stack: Vec, + /// Pending paragraph text buffer. + para_text: String, + /// Heading level for the current paragraph (0 = not a heading). + heading_level: u8, + /// Is the current paragraph a list item? + is_list: bool, + /// Nesting depth of `` elements. + table_depth: u32, + /// Are we inside `` (paragraph properties)? + in_ppr: bool, + /// Are we inside ``? + in_numpr: bool, + /// Are we currently inside a `` text element? + in_text: bool, +} + +impl DocxState { + fn finish_paragraph(&mut self) { + let trimmed = self.para_text.trim(); + if trimmed.is_empty() { + self.para_text.clear(); + self.heading_level = 0; + self.is_list = false; + return; + } + + // Update heading stack. + if self.heading_level > 0 { + // Pop headings at same or deeper level. + self.heading_stack + .truncate((self.heading_level - 1) as usize); + self.heading_stack + .push(format!("Heading {}", self.heading_level)); + } + + let start = self.full_text.len(); + self.full_text.push_str(&self.para_text); + self.full_text.push('\n'); + + let section_kind = if self.table_depth > 0 { + SectionKind::TableCell + } else if self.is_list { + SectionKind::ListItem + } else if self.heading_level > 0 { + SectionKind::Heading + } else { + SectionKind::Paragraph + }; + + self.spans.push(TextSpan { + char_offset_start: offset_to_u32(start), + char_offset_end: offset_to_u32(self.full_text.len()), + heading_path: self.heading_stack.clone(), + section_kind, + }); + + self.para_text.clear(); + self.heading_level = 0; + self.is_list = false; + } +} + +fn parse_docx_xml(xml: &str) -> Result { + let mut reader = Reader::from_str(xml); + reader.config_mut().trim_text(false); + + let mut state = DocxState { + full_text: String::with_capacity(xml.len() / 2), + spans: Vec::new(), + heading_stack: Vec::new(), + para_text: String::new(), + heading_level: 0, + is_list: false, + table_depth: 0, + in_ppr: false, + in_numpr: false, + in_text: false, + }; + + let mut buf = Vec::new(); + loop { + match reader.read_event_into(&mut buf) { + Ok(Event::Start(e)) => match local_name(e.name().as_ref()) { + b"p" => { + state.para_text.clear(); + state.heading_level = 0; + state.is_list = false; + } + b"pPr" => state.in_ppr = true, + b"tbl" => state.table_depth += 1, + b"numPr" => { + state.in_numpr = true; + state.is_list = true; + } + b"t" => state.in_text = true, + _ => {} + }, + Ok(Event::End(e)) => match local_name(e.name().as_ref()) { + b"p" => state.finish_paragraph(), + b"pPr" => state.in_ppr = false, + b"tbl" => state.table_depth = state.table_depth.saturating_sub(1), + b"numPr" => state.in_numpr = false, + b"t" => state.in_text = false, + _ => {} + }, + Ok(Event::Empty(e)) => { + let qname = e.name(); + let ln = local_name(qname.as_ref()); + if ln == b"tab" { + if state.in_text || !state.para_text.is_empty() { + state.para_text.push('\t'); + } + } else if ln == b"br" { + if state.in_text || !state.para_text.is_empty() { + state.para_text.push('\n'); + } + } else if state.in_ppr && ln == b"pStyle" { + if let Some(val) = attr_val(&mut e.attributes(), "val") { + // "Heading1".."Heading6" or "Title" or named styles. + let lower = val.to_ascii_lowercase(); + if let Some(level_str) = lower.strip_prefix("heading") { + if let Ok(level) = level_str.parse::() { + if (1..=6).contains(&level) { + state.heading_level = level; + } + } + } else if lower == "title" { + state.heading_level = 1; + } + } + } + } + Ok(Event::Text(e)) => { + if state.in_text { + let text = decode_xml_entities(&String::from_utf8_lossy(e.as_ref())); + state.para_text.push_str(&text); + } + } + Ok(Event::Eof) => break, + Err(e) => { + return Err(BitVanesError::InvalidInput(format!( + "docx xml parse error at byte {}: {e}", + reader.buffer_position() + ))); + } + _ => {} + } + buf.clear(); + } + + // Guarantee at least one span even for empty docs (chunker expects non-empty + // spans; an empty doc produces zero chunks which is fine). + if state.full_text.is_empty() && state.spans.is_empty() { + // No content — return empty document. + } + + Ok(Document { + full_text: state.full_text, + spans: state.spans, + }) +} + +#[cfg(test)] +#[allow(clippy::needless_raw_string_hashes)] +mod tests { + use super::*; + + /// Builds a minimal valid .docx (ZIP) in memory from the given document.xml body. + fn build_docx(body_xml: &str) -> Vec { + use std::io::Write; + let document_xml = format!( + r#" + + + {body_xml} + +"# + ); + + let mut buf = Vec::new(); + { + let mut writer = zip::ZipWriter::new(std::io::Cursor::new(&mut buf)); + let options = zip::write::SimpleFileOptions::default() + .compression_method(zip::CompressionMethod::Stored); + writer.start_file("word/document.xml", options).unwrap(); + writer.write_all(document_xml.as_bytes()).unwrap(); + writer.finish().unwrap(); + } + buf + } + + #[test] + fn parses_simple_paragraphs() { + let xml = r#" + First paragraph. + Second paragraph. + "#; + let docx = build_docx(xml); + let doc = parse_docx_bytes( + &docx, + &PipelineConfig { + format: DocumentFormat::Docx, + ..PipelineConfig::default() + }, + ) + .unwrap(); + assert!(doc.full_text.contains("First paragraph.")); + assert!(doc.full_text.contains("Second paragraph.")); + assert_eq!(doc.spans.len(), 2); + assert_eq!(doc.spans[0].section_kind, SectionKind::Paragraph); + doc.assert_spans_contiguous(); + } + + #[test] + fn parses_headings_and_body() { + let xml = r#" + My Title + Body text under title. + "#; + let docx = build_docx(xml); + let doc = parse_docx_bytes( + &docx, + &PipelineConfig { + format: DocumentFormat::Docx, + ..PipelineConfig::default() + }, + ) + .unwrap(); + assert_eq!(doc.spans.len(), 2); + assert_eq!(doc.spans[0].section_kind, SectionKind::Heading); + assert!(!doc.spans[0].heading_path.is_empty()); + assert_eq!(doc.spans[1].section_kind, SectionKind::Paragraph); + doc.assert_spans_contiguous(); + } + + #[test] + fn parses_table_cells() { + let xml = r#" + + + Cell A + Cell B + + + "#; + let docx = build_docx(xml); + let doc = parse_docx_bytes( + &docx, + &PipelineConfig { + format: DocumentFormat::Docx, + ..PipelineConfig::default() + }, + ) + .unwrap(); + assert_eq!(doc.spans.len(), 2); + assert!( + doc.spans + .iter() + .all(|s| s.section_kind == SectionKind::TableCell) + ); + doc.assert_spans_contiguous(); + } + + #[test] + fn parses_list_items() { + let xml = r#" + + First item + + Second item + "#; + let docx = build_docx(xml); + let doc = parse_docx_bytes( + &docx, + &PipelineConfig { + format: DocumentFormat::Docx, + ..PipelineConfig::default() + }, + ) + .unwrap(); + assert_eq!(doc.spans.len(), 2); + assert!( + doc.spans + .iter() + .all(|s| s.section_kind == SectionKind::ListItem) + ); + } + + #[test] + fn empty_docx_produces_empty_document() { + let xml = ""; + let docx = build_docx(xml); + let doc = parse_docx_bytes( + &docx, + &PipelineConfig { + format: DocumentFormat::Docx, + ..PipelineConfig::default() + }, + ) + .unwrap(); + assert!(doc.spans.is_empty()); + } + + #[test] + fn invalid_zip_returns_error() { + let result = parse_docx_bytes( + b"not a zip", + &PipelineConfig { + format: DocumentFormat::Docx, + ..PipelineConfig::default() + }, + ); + assert!(result.is_err()); + } + + #[test] + fn xml_entities_are_unescaped() { + let xml = r#"A & B < C"#; + let docx = build_docx(xml); + let doc = parse_docx_bytes( + &docx, + &PipelineConfig { + format: DocumentFormat::Docx, + ..PipelineConfig::default() + }, + ) + .unwrap(); + eprintln!("DEBUG full_text={:?}", doc.full_text); + eprintln!("DEBUG spans={:?}", doc.spans.len()); + assert!( + doc.full_text.contains("A & B < C"), + "text: {:?}", + doc.full_text + ); + assert!(!doc.full_text.contains("&")); + } +} diff --git a/crates/core/src/parse/epub.rs b/crates/core/src/parse/epub.rs new file mode 100644 index 0000000..11560ac --- /dev/null +++ b/crates/core/src/parse/epub.rs @@ -0,0 +1,299 @@ +//! EPUB parser: opens the `.epub` ZIP container, reads the OPF manifest to +//! determine spine (reading) order, then extracts each chapter's XHTML and +//! feeds it through the existing [`HtmlParser`]. +//! +//! This reuses the HTML pipeline — no new parsing logic, just ZIP extraction +//! + OPF ordering. Each chapter's title becomes a heading-ancestry entry. + +use quick_xml::Reader; +use quick_xml::events::Event; + +use crate::error::{BitVanesError, Result}; +use crate::parse::zip_util::read_zip_entry; +use crate::parse::{Document, HtmlParser, Parser, TextSpan, offset_to_u32}; +use crate::schema::{DocumentFormat, PipelineConfig}; + +/// Entry point: parses `.epub` bytes into a [`Document`]. +pub fn parse_epub_bytes(bytes: &[u8], cfg: &PipelineConfig) -> Result { + // 1. Find and parse the OPF file (container.xml → OPF path). + let opf_path = find_opf_path(bytes)?; + let opf_xml = read_zip_entry(bytes, &opf_path)?; + + // 2. Extract spine order + manifest hrefs + titles. + let spine = parse_opf_spine(&opf_xml)?; + + // 3. For each spine item, read its XHTML and parse via HtmlParser. + let html_cfg = PipelineConfig { + format: DocumentFormat::Html, + ..cfg.clone() + }; + let html_parser = HtmlParser; + + let mut full_text = String::new(); + let mut spans: Vec = Vec::new(); + let base_dir = opf_path_dir(&opf_path); + + for item in &spine { + let chapter_path = if item.href.starts_with('/') { + item.href.clone() + } else { + format!("{base_dir}{}", item.href) + }; + + let Ok(xhtml) = read_zip_entry(bytes, &chapter_path) else { + continue; + }; + + let doc = html_parser.parse(&xhtml, &html_cfg)?; + if doc.spans.is_empty() { + continue; + } + + // Merge chapter doc into the main document with a heading prefix. + let heading = vec![item.title.clone()]; + let offset = full_text.len(); + full_text.push_str(&doc.full_text); + + for span in &doc.spans { + let mut combined_heading = heading.clone(); + combined_heading.extend(span.heading_path.iter().cloned()); + spans.push(TextSpan { + char_offset_start: offset_to_u32(offset + span.char_offset_start as usize), + char_offset_end: offset_to_u32(offset + span.char_offset_end as usize), + heading_path: combined_heading, + section_kind: span.section_kind, + }); + } + } + + Ok(Document { full_text, spans }) +} + +struct SpineItem { + href: String, + title: String, +} + +/// Reads `META-INF/container.xml` to find the OPF file path. +fn find_opf_path(bytes: &[u8]) -> Result { + let container = read_zip_entry(bytes, "META-INF/container.xml")?; + let mut reader = Reader::from_str(&container); + let mut buf = Vec::new(); + loop { + match reader.read_event_into(&mut buf) { + Ok(Event::Start(e) | Event::Empty(e)) => { + if local_name(e.name().as_ref()) == b"rootfile" { + for attr in e.attributes().flatten() { + if local_name(attr.key.as_ref()) == b"full-path" { + return Ok(String::from_utf8_lossy(attr.value.as_ref()).into_owned()); + } + } + } + } + Ok(Event::Eof) => break, + Err(e) => { + return Err(BitVanesError::InvalidInput(format!( + "container.xml parse error: {e}" + ))); + } + _ => {} + } + buf.clear(); + } + Err(BitVanesError::InvalidInput( + "container.xml: no element found".to_string(), + )) +} + +/// Parses the OPF `` + `` to produce the reading-order list. +fn parse_opf_spine(opf_xml: &str) -> Result> { + let mut reader = Reader::from_str(opf_xml); + let mut buf = Vec::new(); + + // Collect manifest items: id → href. + let mut manifest: std::collections::HashMap = std::collections::HashMap::new(); + // Collect metadata title. + let mut title = String::from("Chapter#"); + // Spine idrefs in order. + let mut spine_ids: Vec = Vec::new(); + let mut in_metadata = false; + let mut in_title = false; + + loop { + match reader.read_event_into(&mut buf) { + Ok(Event::Start(e)) => match local_name(e.name().as_ref()) { + b"metadata" => in_metadata = true, + b"title" if in_metadata => in_title = true, + b"item" => extract_item(&e, &mut manifest), + b"itemref" => extract_itemref(&e, &mut spine_ids), + _ => {} + }, + // Self-closing elements like and . + Ok(Event::Empty(e)) => match local_name(e.name().as_ref()) { + b"item" => extract_item(&e, &mut manifest), + b"itemref" => extract_itemref(&e, &mut spine_ids), + _ => {} + }, + Ok(Event::End(e)) => match local_name(e.name().as_ref()) { + b"metadata" => in_metadata = false, + b"title" if in_metadata => in_title = false, + _ => {} + }, + Ok(Event::Text(e)) if in_title => { + title = super::docx::decode_xml_entities_pub(&String::from_utf8_lossy(e.as_ref())) + .trim() + .to_string(); + } + Ok(Event::Eof) => break, + Err(e) => return Err(BitVanesError::InvalidInput(format!("opf parse error: {e}"))), + _ => {} + } + buf.clear(); + } + + let mut items = Vec::new(); + let mut chapter_num = 0; + for (i, id) in spine_ids.iter().enumerate() { + if let Some(href) = manifest.get(id) { + chapter_num += 1; + let item_title = if i == 0 && !title.is_empty() { + title.clone() + } else { + format!("Chapter {chapter_num}") + }; + items.push(SpineItem { + href: href.clone(), + title: item_title, + }); + } + } + Ok(items) +} + +/// Returns the directory component of an OPF path (for resolving relative hrefs). +fn opf_path_dir(opf_path: &str) -> String { + match opf_path.rfind('/') { + Some(i) => opf_path[..=i].to_string(), + None => String::new(), + } +} + +/// Extracts `id` + `href` from a `` element. +fn extract_item( + e: &quick_xml::events::BytesStart<'_>, + manifest: &mut std::collections::HashMap, +) { + let mut id = String::new(); + let mut href = String::new(); + for attr in e.attributes().flatten() { + match local_name(attr.key.as_ref()) { + b"id" => id = String::from_utf8_lossy(attr.value.as_ref()).into_owned(), + b"href" => href = String::from_utf8_lossy(attr.value.as_ref()).into_owned(), + _ => {} + } + } + if !id.is_empty() && !href.is_empty() { + manifest.insert(id, href); + } +} + +/// Extracts `idref` from a `` element. +fn extract_itemref(e: &quick_xml::events::BytesStart<'_>, spine_ids: &mut Vec) { + for attr in e.attributes().flatten() { + if local_name(attr.key.as_ref()) == b"idref" { + spine_ids.push(String::from_utf8_lossy(attr.value.as_ref()).into_owned()); + } + } +} + +fn local_name(name: &[u8]) -> &[u8] { + match name.iter().rposition(|&b| b == b':') { + Some(i) => &name[i + 1..], + None => name, + } +} + +#[cfg(test)] +#[allow(clippy::needless_raw_string_hashes)] +mod tests { + use super::*; + use std::io::Write; + + fn build_epub(opf: &str, chapters: &[(&str, &str)]) -> Vec { + let mut buf = Vec::new(); + { + let mut writer = zip::ZipWriter::new(std::io::Cursor::new(&mut buf)); + let opts = zip::write::SimpleFileOptions::default() + .compression_method(zip::CompressionMethod::Stored); + + // container.xml + let container = r#" + + +"#; + writer.start_file("META-INF/container.xml", opts).unwrap(); + writer.write_all(container.as_bytes()).unwrap(); + + // OPF + writer.start_file("OEBPS/content.opf", opts).unwrap(); + writer.write_all(opf.as_bytes()).unwrap(); + + // Chapters + for (path, body) in chapters { + writer.start_file(path, opts).unwrap(); + writer.write_all(body.as_bytes()).unwrap(); + } + writer.finish().unwrap(); + } + buf + } + + #[test] + fn parses_two_chapter_epub() { + let opf = r#" + + + Test Book + + + + + + + + + +"#; + let chapters = [ + ( + "OEBPS/chapter1.xhtml", + "

First chapter text.

", + ), + ( + "OEBPS/chapter2.xhtml", + "

Second chapter text.

", + ), + ]; + let epub = build_epub(opf, &chapters); + let doc = parse_epub_bytes( + &epub, + &PipelineConfig { + format: DocumentFormat::Epub, + ..PipelineConfig::default() + }, + ) + .unwrap(); + + assert!(doc.full_text.contains("First chapter text.")); + assert!(doc.full_text.contains("Second chapter text.")); + assert!(doc.spans.len() >= 2); + // First chapter heading should carry the book title. + assert!( + doc.spans[0] + .heading_path + .iter() + .any(|h| h.contains("Test Book")) + ); + doc.assert_spans_contiguous(); + } +} diff --git a/crates/core/src/parse/pptx.rs b/crates/core/src/parse/pptx.rs new file mode 100644 index 0000000..5f8c2d8 --- /dev/null +++ b/crates/core/src/parse/pptx.rs @@ -0,0 +1,173 @@ +//! PPTX (Microsoft `PowerPoint`) parser via ZIP + quick-xml. +//! +//! Opens the `.pptx` ZIP container, enumerates `ppt/slides/slideN.xml` files +//! in numeric order, and extracts text runs (``) from each slide. Each +//! slide becomes a [`TextSpan`] with `heading_path = ["Slide N"]`, so chunks +//! carry slide-level lineage for RAG retrieval. + +use quick_xml::Reader; +use quick_xml::events::Event; + +use crate::error::{BitVanesError, Result}; +use crate::parse::zip_util::{list_entries, read_zip_entry}; +use crate::parse::{Document, TextSpan, offset_to_u32}; +use crate::schema::{PipelineConfig, SectionKind}; + +/// Entry point: parses `.pptx` bytes into a [`Document`]. +pub fn parse_pptx_bytes(bytes: &[u8], _cfg: &PipelineConfig) -> Result { + let mut slide_names = list_entries(bytes, "ppt/slides/")?; + // Sort by slide number (slide1.xml < slide2.xml < ... slide10.xml). + slide_names.sort_by(|a, b| { + let na = slide_number(a); + let nb = slide_number(b); + na.cmp(&nb) + }); + slide_names.retain(|n| n.to_ascii_lowercase().ends_with(".xml")); + + let mut full_text = String::new(); + let mut spans: Vec = Vec::new(); + + for name in &slide_names { + let xml = read_zip_entry(bytes, name)?; + let xml = super::docx::encode_xml_entities_pub(&xml); + let slide_num = slide_number(name).unwrap_or(0); + let heading = format!("Slide {slide_num}"); + + let text = extract_slide_text(&xml)?; + if text.trim().is_empty() { + continue; + } + + let start = full_text.len(); + full_text.push_str(&text); + full_text.push('\n'); + + spans.push(TextSpan { + char_offset_start: offset_to_u32(start), + char_offset_end: offset_to_u32(full_text.len()), + heading_path: vec![heading], + section_kind: SectionKind::Paragraph, + }); + } + + Ok(Document { full_text, spans }) +} + +/// Extracts the numeric suffix from a slide path like `ppt/slides/slide3.xml`. +fn slide_number(name: &str) -> Option { + let base = name.rsplit('/').next()?; + let num_part = base.strip_prefix("slide")?; + let num_part = num_part.strip_suffix(".xml")?; + num_part.parse().ok() +} + +/// Walks a slide's XML and collects all `` text content, joined by newlines. +fn extract_slide_text(xml: &str) -> Result { + let mut reader = Reader::from_str(xml); + reader.config_mut().trim_text(false); + + let mut text = String::new(); + let mut in_t = false; + let mut buf = Vec::new(); + + loop { + match reader.read_event_into(&mut buf) { + Ok(Event::Start(e)) if local_name(e.name().as_ref()) == b"t" => { + in_t = true; + if !text.is_empty() && !text.ends_with('\n') { + text.push('\n'); + } + } + Ok(Event::End(e)) if local_name(e.name().as_ref()) == b"t" => in_t = false, + Ok(Event::Text(e)) if in_t => { + let decoded = + super::docx::decode_xml_entities_pub(&String::from_utf8_lossy(e.as_ref())); + text.push_str(&decoded); + } + Ok(Event::Eof) => break, + Err(e) => { + return Err(BitVanesError::InvalidInput(format!( + "pptx xml parse error: {e}" + ))); + } + _ => {} + } + buf.clear(); + } + Ok(text) +} + +fn local_name(name: &[u8]) -> &[u8] { + match name.iter().rposition(|&b| b == b':') { + Some(i) => &name[i + 1..], + None => name, + } +} + +#[cfg(test)] +#[allow(clippy::needless_raw_string_hashes)] +mod tests { + use super::*; + use crate::schema::DocumentFormat; + use std::io::Write; + + fn build_pptx(slides: &[(usize, &str)]) -> Vec { + let mut buf = Vec::new(); + { + let mut writer = zip::ZipWriter::new(std::io::Cursor::new(&mut buf)); + let opts = zip::write::SimpleFileOptions::default() + .compression_method(zip::CompressionMethod::Stored); + for (num, body) in slides { + let path = format!("ppt/slides/slide{num}.xml"); + writer.start_file(&path, opts).unwrap(); + let xml = format!( + r#" + {body} + "# + ); + writer.write_all(xml.as_bytes()).unwrap(); + } + writer.finish().unwrap(); + } + buf + } + + #[test] + fn parses_multiple_slides_in_order() { + let pptx = build_pptx(&[ + (1, r#"WelcomeIntro slide"#), + (2, r#"Data results"#), + (10, r#"Tenth slide"#), + ]); + let doc = parse_pptx_bytes( + &pptx, + &PipelineConfig { + format: DocumentFormat::Pptx, + ..PipelineConfig::default() + }, + ) + .unwrap(); + assert_eq!(doc.spans.len(), 3); + assert!(doc.full_text.contains("Welcome")); + assert!(doc.full_text.contains("Tenth slide")); + // Slides sorted numerically: slide 1, 2, 10 (not lexicographic). + assert_eq!(doc.spans[0].heading_path, vec!["Slide 1"]); + assert_eq!(doc.spans[1].heading_path, vec!["Slide 2"]); + assert_eq!(doc.spans[2].heading_path, vec!["Slide 10"]); + doc.assert_spans_contiguous(); + } + + #[test] + fn empty_slides_skipped() { + let pptx = build_pptx(&[(1, r#""#), (2, r#"Content"#)]); + let doc = parse_pptx_bytes( + &pptx, + &PipelineConfig { + format: DocumentFormat::Pptx, + ..PipelineConfig::default() + }, + ) + .unwrap(); + assert_eq!(doc.spans.len(), 1); + } +} diff --git a/crates/core/src/parse/rtf.rs b/crates/core/src/parse/rtf.rs new file mode 100644 index 0000000..24e0c2f --- /dev/null +++ b/crates/core/src/parse/rtf.rs @@ -0,0 +1,97 @@ +//! RTF (Rich Text Format) parser via `rtf-parser`. +//! +//! Extracts plain text + paragraph structure from RTF documents. Each +//! `StyleBlock` in the document body becomes a candidate span; empty blocks +//! are collapsed. Heading detection is heuristic (ALL-CAPS short lines or +//! bold painter) since RTF style sheets are inconsistently applied. + +use rtf_parser::parse_rtf; + +use crate::error::{BitVanesError, Result}; +use crate::parse::{Document, TextSpan, offset_to_u32}; +use crate::schema::{PipelineConfig, SectionKind}; + +/// Entry point: parses RTF bytes into a [`Document`]. +pub fn parse_rtf_bytes(bytes: &[u8], _cfg: &PipelineConfig) -> Result { + let text = std::str::from_utf8(bytes) + .map_err(|e| BitVanesError::InvalidInput(format!("rtf is not valid UTF-8/ASCII: {e}")))?; + + let rtf_doc = parse_rtf(text.to_string()); + + let mut full_text = String::new(); + let mut spans: Vec = Vec::new(); + let mut heading_stack: Vec = Vec::new(); + + for block in &rtf_doc.body { + let para_text = block.text.trim(); + if para_text.is_empty() { + continue; + } + + // Heuristic heading: short ALL-CAPS line. + let alpha_count = para_text.chars().filter(|c| c.is_alphabetic()).count(); + let is_heading = para_text.len() < 80 + && alpha_count >= 4 + && para_text + .chars() + .filter(|c| c.is_alphabetic()) + .all(char::is_uppercase); + + if is_heading { + heading_stack.clear(); + heading_stack.push(para_text.to_string()); + } + + let start = full_text.len(); + full_text.push_str(para_text); + full_text.push('\n'); + + spans.push(TextSpan { + char_offset_start: offset_to_u32(start), + char_offset_end: offset_to_u32(full_text.len()), + heading_path: if is_heading { + Vec::new() + } else { + heading_stack.clone() + }, + section_kind: if is_heading { + SectionKind::Heading + } else { + SectionKind::Paragraph + }, + }); + } + + Ok(Document { full_text, spans }) +} + +#[cfg(test)] +#[allow(clippy::needless_raw_string_hashes)] +mod tests { + use super::*; + use crate::schema::DocumentFormat; + + #[test] + fn parses_simple_rtf() { + let rtf = r#"{\rtf1\ansi\deff0 +{\fonttbl{\f0 Times New Roman;}} +\f0\fs24 +Hello world.\par +This is a second paragraph.\par +}"#; + let doc = parse_rtf_bytes( + rtf.as_bytes(), + &PipelineConfig { + format: DocumentFormat::Rtf, + ..PipelineConfig::default() + }, + ) + .unwrap(); + assert!( + doc.full_text.contains("Hello world") || !doc.spans.is_empty(), + "should extract text: {:?}", + doc.full_text + ); + assert!(!doc.spans.is_empty()); + } +} diff --git a/crates/core/src/parse/xlsx.rs b/crates/core/src/parse/xlsx.rs new file mode 100644 index 0000000..084e185 --- /dev/null +++ b/crates/core/src/parse/xlsx.rs @@ -0,0 +1,174 @@ +//! XLSX (Microsoft Excel) parser via `calamine`. +//! +//! # Memory strategy for large/sparse spreadsheets +//! +//! Calamine materialises each sheet into a dense `Range` grid. +//! For a sheet with dimension 1,000,000 × 50 but only 200 populated cells, +//! the full grid is still allocated. To prevent OOM on pathological files: +//! +//! 1. **One sheet at a time**: `worksheet_range(name)` loads a single sheet, +//! not `worksheets()` which loads all sheets simultaneously. +//! 2. **Row iterator with empty-row trimming**: `range.rows()` gives a +//! slice view; we skip all-empty rows and break at the first long run of +//! empty trailing rows (calamine's used-range bounding box usually +//! handles this, but we belt-and-braces it). +//! 3. **Cell-count guard**: `max_cells` (default 2,000,000) caps total +//! materialised cells to prevent unbounded allocation. If exceeded, the +//! sheet is truncated and a warning is emitted. +//! 4. **Sparse cell skipping**: `DataType::Empty` cells are not emitted as +//! text — they're collapsed into a single tab delimiter. +//! +//! Each contiguous block of non-empty rows becomes one span. The sheet name +//! becomes the heading ancestry: `["Sheet1"]`. + +use calamine::{Data, Reader as CalamineReader, open_workbook_auto_from_rs}; + +use crate::error::{BitVanesError, Result}; +use crate::parse::{Document, TextSpan, offset_to_u32}; +use crate::schema::{PipelineConfig, SectionKind}; + +/// Maximum cells to materialise per workbook before truncating. At ~50 bytes +/// per `Data` enum value, 2M cells ≈ 100 MB — a reasonable ceiling. +const MAX_CELLS: usize = 2_000_000; + +/// Maximum contiguous empty rows before we stop scanning a sheet (trailing +/// junk that calamine's bounding box sometimes includes). +const MAX_TRAILING_EMPTY: usize = 100; + +/// Entry point: parses `.xlsx` bytes into a [`Document`]. +pub fn parse_xlsx_bytes(bytes: &[u8], _cfg: &PipelineConfig) -> Result { + let cursor = std::io::Cursor::new(bytes); + let mut workbook = open_workbook_auto_from_rs(cursor) + .map_err(|e| BitVanesError::InvalidInput(format!("xlsx open error: {e}")))?; + + let sheet_names = workbook.sheet_names().clone(); + + let mut full_text = String::new(); + let mut spans: Vec = Vec::new(); + let mut total_cells = 0usize; + + for sheet_name in &sheet_names { + if total_cells >= MAX_CELLS { + break; + } + + let Ok(range) = workbook.worksheet_range(sheet_name) else { + continue; + }; + + let height = range.height(); + let width = range.width().max(1); + // Guard: skip sheets that would blow the cell budget. + if total_cells.saturating_add(height * width) > MAX_CELLS { + break; + } + total_cells += height * width; + + let heading = vec![sheet_name.clone()]; + let mut row_block = String::new(); + let mut empty_streak = 0u32; + + for row in range.rows() { + let non_empty: Vec<&Data> = row.iter().filter(|c| !matches!(c, Data::Empty)).collect(); + + if non_empty.is_empty() { + empty_streak += 1; + // Flush the current row block as a span. + if !row_block.trim().is_empty() { + push_span(&mut full_text, &mut spans, &row_block, &heading); + row_block.clear(); + } + if empty_streak >= u32::try_from(MAX_TRAILING_EMPTY).unwrap_or(u32::MAX) { + break; + } + continue; + } + empty_streak = 0; + + // Build a tab-separated row string, collapsing empty cells. + let row_str: String = row + .iter() + .map(cell_to_string) + .collect::>() + .join("\t"); + + if !row_block.is_empty() { + row_block.push('\n'); + } + row_block.push_str(&row_str); + } + + if !row_block.trim().is_empty() { + push_span(&mut full_text, &mut spans, &row_block, &heading); + } + } + + Ok(Document { full_text, spans }) +} + +/// Converts a calamine `Data` cell to its string representation. +/// `Data::Empty` → empty string (collapsed by the join delimiter). +#[allow(clippy::cast_possible_truncation)] +fn cell_to_string(cell: &Data) -> String { + match cell { + Data::Empty => String::new(), + Data::String(s) | Data::DurationIso(s) | Data::DateTimeIso(s) => s.clone(), + Data::Int(n) => n.to_string(), + Data::Float(f) => { + // Avoid trailing ".0" for whole numbers. + if f.fract() == 0.0 && f.is_finite() { + format!("{f:.0}") + } else { + f.to_string() + } + } + Data::DateTime(dt) => dt.to_string(), + Data::Bool(b) => b.to_string(), + Data::Error(err) => format!("#ERROR:{err:?}"), + } +} + +/// Appends a text block to the full document and records a span. +fn push_span(full_text: &mut String, spans: &mut Vec, text: &str, heading: &[String]) { + let start = full_text.len(); + full_text.push_str(text); + full_text.push('\n'); + spans.push(TextSpan { + char_offset_start: offset_to_u32(start), + char_offset_end: offset_to_u32(full_text.len()), + heading_path: heading.to_vec(), + section_kind: SectionKind::TableCell, + }); +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::schema::DocumentFormat; + use calamine::Data; + + /// Tests cell conversion logic and error handling. + /// Full XLSX round-trip tests live in `tests/formats.rs` with fixture files. + + #[test] + fn cell_to_string_handles_types() { + assert_eq!(cell_to_string(&Data::Empty), ""); + assert_eq!(cell_to_string(&Data::Int(42)), "42"); + assert_eq!(cell_to_string(&Data::Float(3.15)), "3.15"); + assert_eq!(cell_to_string(&Data::Float(100.0)), "100"); + assert_eq!(cell_to_string(&Data::String("hello".into())), "hello"); + assert_eq!(cell_to_string(&Data::Bool(true)), "true"); + } + + #[test] + fn invalid_xlsx_returns_error() { + let result = parse_xlsx_bytes( + b"not an xlsx", + &PipelineConfig { + format: DocumentFormat::Xlsx, + ..PipelineConfig::default() + }, + ); + assert!(result.is_err()); + } +} diff --git a/crates/core/src/parse/zip_util.rs b/crates/core/src/parse/zip_util.rs new file mode 100644 index 0000000..ea70697 --- /dev/null +++ b/crates/core/src/parse/zip_util.rs @@ -0,0 +1,63 @@ +//! Shared helpers for ZIP-based document formats (DOCX, PPTX, EPUB). +//! +//! All three are ZIP containers holding XML/XHTML parts. This module +//! provides the common `read_zip_entry` function so each parser only +//! implements its format-specific XML walking logic. + +use std::io::Read; + +use zip::ZipArchive; + +use crate::error::{BitVanesError, Result}; + +/// Opens a ZIP archive from raw bytes and returns the decoded UTF-8 content +/// of the named entry. +/// +/// # Errors +/// +/// Returns [`BitVanesError::InvalidInput`] if the bytes are not a valid ZIP +/// archive, the entry is not found, or the entry is not valid UTF-8. +pub fn read_zip_entry(bytes: &[u8], entry_name: &str) -> Result { + let cursor = std::io::Cursor::new(bytes); + let mut archive = ZipArchive::new(cursor) + .map_err(|e| BitVanesError::InvalidInput(format!("invalid zip archive: {e}")))?; + let mut file = archive + .by_name(entry_name) + .map_err(|e| BitVanesError::InvalidInput(format!("zip entry '{entry_name}': {e}")))?; + let mut buf = String::new(); + file.read_to_string(&mut buf) + .map_err(|e| BitVanesError::InvalidInput(format!("zip read error: {e}")))?; + Ok(buf) +} + +/// Returns the raw bytes of a ZIP entry (for binary entries like images). +pub fn read_zip_entry_bytes(bytes: &[u8], entry_name: &str) -> Result> { + let cursor = std::io::Cursor::new(bytes); + let mut archive = ZipArchive::new(cursor) + .map_err(|e| BitVanesError::InvalidInput(format!("invalid zip archive: {e}")))?; + let mut file = archive + .by_name(entry_name) + .map_err(|e| BitVanesError::InvalidInput(format!("zip entry '{entry_name}': {e}")))?; + let mut buf = Vec::new(); + file.read_to_end(&mut buf) + .map_err(|e| BitVanesError::InvalidInput(format!("zip read error: {e}")))?; + Ok(buf) +} + +/// Enumerates all entry names in the archive that start with `prefix`. +pub fn list_entries(bytes: &[u8], prefix: &str) -> Result> { + let cursor = std::io::Cursor::new(bytes); + let mut archive = ZipArchive::new(cursor) + .map_err(|e| BitVanesError::InvalidInput(format!("invalid zip archive: {e}")))?; + let mut names = Vec::new(); + for i in 0..archive.len() { + let entry = archive + .by_index(i) + .map_err(|e| BitVanesError::InvalidInput(format!("zip index error: {e}")))?; + let name = entry.name().to_string(); + if name.starts_with(prefix) { + names.push(name); + } + } + Ok(names) +} diff --git a/crates/core/src/pii_detect.rs b/crates/core/src/pii_detect.rs new file mode 100644 index 0000000..390f2ef --- /dev/null +++ b/crates/core/src/pii_detect.rs @@ -0,0 +1,143 @@ +//! Tier-2 pluggable PII detection: an extension point for ML/NLP-based +//! named-entity recognisers (NER) that complement the Tier-1 regex + +//! checksum engine in [`crate::scrub`]. +//! +//! # Design +//! +//! The Tier-1 engine (regex + Luhn/ABA + anchor windows) is fast, +//! deterministic, and has no external dependencies. Some PII — personal +//! names, physical addresses, dates of birth, medical record numbers — +//! resists structural patterns and requires a statistical model. The +//! [`PiiDetector`] trait is the plug-in point for such models. +//! +//! A detector runs **after** Tier-1 and appends its findings to the same +//! `Vec`. Overlap resolution then merges Tier-1 and Tier-2 +//! detections. This keeps the pipeline contract unchanged: one finding list, +//! one offset map, one `pii_metadata` column. +//! +//! # Current status +//! +//! [`ModelDetector`] is a **stub**: it exists so the trait object can be +//! wired through the pipeline and config today, but its `detect` impl is +//! `unimplemented!()` behind the `pii-model` cargo feature. A future pass +//! will provide an ONNX-backed implementation (BERT-NER or similar). +//! +//! # Future: ONNX integration +//! +//! ```ignore +//! let detector = ModelDetector::new("models/bert-ner.onnx", "models/tokenizer.json")?; +//! // detector.detect(&text, &mut findings)?; +//! ``` + +use crate::error::Result; +use crate::scrub::PiiFinding; + +/// A pluggable Tier-2 PII detector (NER model, gazetteer, ...). +/// +/// Implementations append findings to the caller-provided `Vec`; the +/// pipeline merges them with Tier-1 results and resolves overlaps. +/// +/// `Send + Sync` is required so the detector can run on a rayon thread pool +/// in the CLI's batch path. +pub trait PiiDetector: Send + Sync { + /// Scans `text` and appends any detected findings to `findings`. + /// + /// Findings MUST carry offsets into the same `text` (original-text + /// coordinate space, before scrubbing). + /// + /// # Errors + /// + /// Implementations may return [`crate::error::BitVanesError`] on model + /// load failure, inference error, or invalid input. + fn detect(&self, text: &str, findings: &mut Vec) -> Result<()>; +} + +/// Placeholder Tier-2 detector backed by an (unimplemented) ONNX NER model. +/// +/// Fields store the model + tokenizer paths so a future implementation can +/// lazy-load them. Constructed via [`ModelDetector::new`]. +/// +/// **Not yet functional** — `detect()` panics with `unimplemented!()`. +/// Gated behind the `pii-model` cargo feature. +#[derive(Debug)] +#[allow(dead_code)] +pub struct ModelDetector { + model_path: String, + tokenizer_path: String, + /// Confidence floor for model-emitted findings. Model logits are often + /// noisy; findings below this floor are dropped. + min_confidence: f32, +} + +impl ModelDetector { + /// Creates a detector handle for the given model + tokenizer paths. + /// + /// Does NOT load the model (lazy-load happens on first `detect` call in + /// the future implementation). + #[must_use] + pub fn new(model_path: impl Into, tokenizer_path: impl Into) -> Self { + Self { + model_path: model_path.into(), + tokenizer_path: tokenizer_path.into(), + min_confidence: 0.70, + } + } + + /// Sets the minimum confidence for model-emitted findings. + #[must_use] + pub const fn with_min_confidence(mut self, floor: f32) -> Self { + self.min_confidence = floor; + self + } +} + +#[cfg(feature = "pii-model")] +impl PiiDetector for ModelDetector { + fn detect(&self, _text: &str, _findings: &mut Vec) -> Result<()> { + // TODO(phase-6): load ONNX model + tokenizer, run inference, map + // token-level NER labels back to character offsets, filter by + // min_confidence, and append PiiFindings. + unimplemented!("ModelDetector is a stub; ONNX integration lands in a future phase") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn model_detector_stores_paths() { + let d = ModelDetector::new("a.onnx", "b.json"); + assert_eq!(d.model_path, "a.onnx"); + assert_eq!(d.tokenizer_path, "b.json"); + assert!((d.min_confidence - 0.70).abs() < 1e-6); + } + + #[test] + fn model_detector_builder_sets_confidence() { + let d = ModelDetector::new("a.onnx", "b.json").with_min_confidence(0.85); + assert!((d.min_confidence - 0.85).abs() < 1e-6); + } + + #[test] + fn noop_detector_impl_works() { + // A trivial no-op detector proves the trait is object-safe and + // composable with the pipeline without needing the ONNX feature. + struct Noop; + impl PiiDetector for Noop { + fn detect(&self, _text: &str, findings: &mut Vec) -> Result<()> { + findings.clear(); // explicitly empty + Ok(()) + } + } + let mut v = vec![PiiFinding { + entity: "email".to_string(), + offset_start: 0, + offset_end: 5, + confidence: 0.9, + anchors_hit: vec![], + }]; + Noop.detect("hello", &mut v).unwrap(); + assert!(v.is_empty()); + } +} diff --git a/crates/core/src/pipeline.rs b/crates/core/src/pipeline.rs index 7e00882..d9b4509 100644 --- a/crates/core/src/pipeline.rs +++ b/crates/core/src/pipeline.rs @@ -7,29 +7,31 @@ use arrow::array::RecordBatch; use crate::arrow_io::batch::{chunks_to_batch, chunks_to_batch_with_embeddings}; -use crate::chunk::chunk_document; +use crate::chunk::{chunk_document, chunk_document_semantic}; use crate::embed::Embedder; use crate::error::Result; use crate::parse::parse_bytes; -use crate::schema::PipelineConfig; -use crate::scrub::scrub_document; +use crate::schema::{ChunkStrategy, PipelineConfig}; +use crate::scrub::{OffsetMap, PiiFinding, scrub_document}; /// Runs the full ETL pipeline on `bytes` and returns an Arrow /// [`RecordBatch`] containing the chunked output. /// /// Stages: /// 1. Parse raw bytes into a [`Document`] via the configured format. -/// 2. Scrub PII via the configured [`ScrubProfile`]. +/// 2. Scrub PII via the configured [`ScrubProfile`] (emits findings + offsets). /// 3. Chunk the scrubbed text via BPE token boundaries. -/// 4. Assemble chunks into an Arrow `RecordBatch`. +/// 4. Attach per-chunk PII findings + deterministic `chunk_id`. +/// 5. Assemble chunks into an Arrow `RecordBatch`. /// /// # Errors /// /// Propagates [`crate::error::BitVanesError`] from any stage. pub fn run_pipeline(bytes: &[u8], cfg: &PipelineConfig) -> Result { let doc = parse_bytes(bytes, cfg)?; - let (scrubbed_doc, _offset_map) = scrub_document(doc, &cfg.scrub)?; - let chunks = chunk_document(&scrubbed_doc, &cfg.chunk, cfg.source_label.as_deref())?; + let (scrubbed_doc, offset_map, findings) = scrub_document(doc, &cfg.scrub)?; + let mut chunks = chunk_document(&scrubbed_doc, &cfg.chunk, cfg.source_label.as_deref())?; + attach_metadata(&mut chunks, &findings, &offset_map); let batch = chunks_to_batch(&chunks)?; Ok(batch) } @@ -52,8 +54,9 @@ pub fn run_pipeline_with_embeddings( embedder: &dyn Embedder, ) -> Result { let doc = parse_bytes(bytes, cfg)?; - let (scrubbed_doc, _offset_map) = scrub_document(doc, &cfg.scrub)?; - let chunks = chunk_document(&scrubbed_doc, &cfg.chunk, cfg.source_label.as_deref())?; + let (scrubbed_doc, offset_map, findings) = scrub_document(doc, &cfg.scrub)?; + let mut chunks = chunk_document(&scrubbed_doc, &cfg.chunk, cfg.source_label.as_deref())?; + attach_metadata(&mut chunks, &findings, &offset_map); let texts: Vec<&str> = chunks.iter().map(|c| c.text.as_str()).collect(); let embeddings = embedder.embed(&texts)?; @@ -63,6 +66,37 @@ pub fn run_pipeline_with_embeddings( Ok(batch) } +/// Like [`run_pipeline`] but honours [`ChunkStrategy::Semantic`]: when the +/// config asks for semantic chunking, `embedder` guides where cuts happen. +/// For [`ChunkStrategy::Structural`] the embedder is unused and this is +/// equivalent to [`run_pipeline`]. +/// +/// # Errors +/// +/// Propagates [`crate::error::BitVanesError`] from any stage or the embedder. +pub fn run_pipeline_with_strategy( + bytes: &[u8], + cfg: &PipelineConfig, + embedder: &dyn Embedder, +) -> Result { + let doc = parse_bytes(bytes, cfg)?; + let (scrubbed_doc, offset_map, findings) = scrub_document(doc, &cfg.scrub)?; + let mut chunks = match cfg.chunk.strategy { + ChunkStrategy::Structural => { + chunk_document(&scrubbed_doc, &cfg.chunk, cfg.source_label.as_deref())? + } + ChunkStrategy::Semantic { .. } => chunk_document_semantic( + &scrubbed_doc, + &cfg.chunk, + embedder, + cfg.source_label.as_deref(), + )?, + }; + attach_metadata(&mut chunks, &findings, &offset_map); + let batch = chunks_to_batch(&chunks)?; + Ok(batch) +} + /// Runs [`run_pipeline`] over many inputs, in parallel when the `parallel` /// feature is enabled and sequentially otherwise. Returns one result per /// input so a single failing document does not abort the batch. @@ -83,6 +117,43 @@ pub fn run_pipeline_batch(inputs: &[&[u8]], cfg: &PipelineConfig) -> Vec orig_start + }) + .cloned() + .collect(); + + // Deterministic content-addressed ID: blake3(text || source || offsets). + let mut hasher = blake3::Hasher::new(); + hasher.update(chunk.text.as_bytes()); + hasher.update(chunk.source_path.as_bytes()); + hasher.update(&chunk.char_offset_start.to_le_bytes()); + hasher.update(&chunk.char_offset_end.to_le_bytes()); + chunk.chunk_id = hasher.finalize().to_hex().to_string(); + } +} + #[cfg(test)] mod batch_tests { use super::*; @@ -123,7 +194,9 @@ mod batch_tests { #[cfg(test)] mod tests { use super::*; - use crate::schema::{ChunkConfig, DocumentFormat, PipelineConfig, ScrubProfile}; + use crate::schema::{ + BuiltInPattern, ChunkConfig, DocumentFormat, PipelineConfig, ScrubProfile, + }; #[test] fn pipeline_produces_nonempty_batch_from_markdown() { @@ -140,7 +213,7 @@ mod tests { let input = b"# Title\n\nHello world. This is a test."; let batch = run_pipeline(input, &cfg).unwrap(); assert!(batch.num_rows() > 0); - assert_eq!(batch.num_columns(), 9); + assert_eq!(batch.num_columns(), 11); } #[test] @@ -150,8 +223,8 @@ mod tests { let cfg = PipelineConfig { format: DocumentFormat::Markdown, scrub: ScrubProfile { - patterns: vec![crate::schema::BuiltInPattern::Email], - custom: vec![], + patterns: vec![BuiltInPattern::Email], + ..ScrubProfile::default() }, chunk: ChunkConfig::default(), source_label: None, @@ -162,7 +235,7 @@ mod tests { assert_eq!(batch.num_rows(), 1); let text_col = batch - .column(1) + .column(2) .as_any() .downcast_ref::() .unwrap(); @@ -176,6 +249,67 @@ mod tests { ); } + #[test] + fn pipeline_emits_pii_metadata_for_scrubbed_email() { + use arrow::array::{Array, ListArray}; + + let cfg = PipelineConfig { + format: DocumentFormat::Markdown, + scrub: ScrubProfile { + patterns: vec![BuiltInPattern::Email], + ..ScrubProfile::default() + }, + chunk: ChunkConfig::default(), + source_label: None, + embeddings: None, + }; + let input = b"Contact alice@example.com for info."; + let batch = run_pipeline(input, &cfg).unwrap(); + + let pii_col = batch + .column(9) + .as_any() + .downcast_ref::() + .expect("pii_metadata should be List"); + assert!(!pii_col.is_null(0), "row 0 should carry a finding"); + } + + #[test] + fn pipeline_chunk_id_is_deterministic() { + use arrow::array::{Array, StringArray}; + + let cfg = PipelineConfig { + format: DocumentFormat::Markdown, + scrub: ScrubProfile { + patterns: vec![BuiltInPattern::Email], + ..ScrubProfile::default() + }, + chunk: ChunkConfig::default(), + source_label: None, + embeddings: None, + }; + let input = b"Contact alice@example.com for info."; + let batch1 = run_pipeline(input, &cfg).unwrap(); + let batch2 = run_pipeline(input, &cfg).unwrap(); + + let id_col1 = batch1 + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let id_col2 = batch2 + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + id_col1.value(0), + id_col2.value(0), + "chunk_id must be stable" + ); + assert_eq!(id_col1.value(0).len(), 64, "blake3 hex is 64 chars"); + } + #[test] fn pipeline_empty_input_produces_empty_batch() { let cfg = PipelineConfig::default(); @@ -196,12 +330,10 @@ mod tests { assert!(batch.num_rows() >= 1); let heading_col = batch - .column(4) + .column(5) .as_any() .downcast_ref::() .unwrap(); - // The chunk should have a non-null heading_path (under "Architecture" - // or "Architecture > Storage"). let has_heading = (0..heading_col.len()).any(|i| !heading_col.is_null(i)); assert!(has_heading, "at least one chunk should have heading_path"); } diff --git a/crates/core/src/schema.rs b/crates/core/src/schema.rs index fe5ac17..abbe459 100644 --- a/crates/core/src/schema.rs +++ b/crates/core/src/schema.rs @@ -25,7 +25,7 @@ use serde::{Deserialize, Serialize}; /// Identical JSON shape is consumed by both the wasm web worker and the /// native CLI - this is the wire format of the Milestone 4 `profile.json` /// exported by the visual studio. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)] pub struct PipelineConfig { /// Source document format. Selects the parser. Defaults to /// [`DocumentFormat::Markdown`] via the derived `Default` impl. @@ -87,16 +87,17 @@ pub struct EmbeddingConfig { /// Supported source document formats. /// -/// # Note on PDF +/// # Binary formats (DOCX, PPTX, XLSX, EPUB, RTF, PDF) /// -/// `Pdf` is only resolvable on native targets with the `cli-pdf` feature -/// enabled. The web PDF path runs through Mozilla PDF.js in a dedicated -/// JavaScript worker *before* reaching this engine, so the wasm build -/// receives the extracted text as [`Markdown`] or [`Text`]. +/// These are binary container formats (ZIP or proprietary) that require +/// native feature-gated parsers. They are routed by [`parse_bytes`] before +/// UTF-8 decoding. The wasm build cannot parse them — the browser must +/// extract text via JS libraries (PDF.js, mammoth.js, etc.) and pass the +/// resulting text/markdown to the engine. /// -/// [`Markdown`]: DocumentFormat::Markdown -/// [`Text`]: DocumentFormat::Text -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +/// [`parse_bytes`]: crate::parse::parse_bytes +#[non_exhaustive] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)] #[serde(rename_all = "lowercase")] pub enum DocumentFormat { /// GitHub-flavored Markdown via `pulldown-cmark`. @@ -111,6 +112,16 @@ pub enum DocumentFormat { Html, /// PDF (native only; requires the `cli-pdf` feature). Pdf, + /// Microsoft Word `.docx` (native only; requires the `office` feature). + Docx, + /// Microsoft `PowerPoint` `.pptx` (native only; requires `office`). + Pptx, + /// Microsoft Excel `.xlsx` (native only; requires `office`). + Xlsx, + /// EPUB ebook (native only; requires `office`). + Epub, + /// Rich Text Format `.rtf` (native only; requires `office`). + Rtf, } /// Which BPE tokenizer to apply when computing chunk boundaries. @@ -136,9 +147,34 @@ pub enum TokenizerKind { O200kHarmony, } +/// How the chunker decides where to cut. +/// +/// `Structural` packs spans greedily up to `max_tokens` at structural +/// boundaries. `Semantic` additionally requires that consecutive spans stay +/// above a cosine-similarity threshold so a chunk holds a single topical +/// unit; cutting happens when the topic drifts *or* `max_tokens` is hit. +/// +/// Note: this enum carries an `f32` threshold, so `ChunkConfig` and +/// `PipelineConfig` implement `PartialEq` but not `Eq` as of 0.2.0. +#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, Default)] +#[serde(rename_all = "lowercase")] +pub enum ChunkStrategy { + /// Greedy structural packing up to `max_tokens` (the original behaviour). + #[default] + Structural, + /// Embedding-guided merging. Adjacent spans are merged into the current + /// chunk while their cosine similarity to the running chunk centroid is + /// `>= similarity_threshold` *and* `max_tokens` is respected. + Semantic { + /// Minimum cosine similarity between a candidate span and the current + /// chunk's centroid to keep merging. Typical range 0.5–0.85. + similarity_threshold: f32, + }, +} + /// Chunking parameters. All fields are validated by the engine at /// pipeline-start time; see [`crate::chunk`] for the enforcement site. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct ChunkConfig { /// Target maximum token count per chunk. The chunker saturates up to /// but never exceeds this value. Must be greater than zero. @@ -146,12 +182,18 @@ pub struct ChunkConfig { /// Number of tokens of overlap between adjacent chunks. Defaults to /// zero (no overlap). Must be strictly less than `max_tokens`. + /// Overlap is honoured by [`ChunkStrategy::Structural`]; the semantic + /// strategy ignores it (treat as 0). #[serde(default)] pub overlap_tokens: u32, /// Which BPE tokenizer to apply. #[serde(default)] pub tokenizer: TokenizerKind, + + /// How chunk boundaries are chosen. Defaults to structural packing. + #[serde(default)] + pub strategy: ChunkStrategy, } impl Default for ChunkConfig { @@ -160,6 +202,7 @@ impl Default for ChunkConfig { max_tokens: 512, overlap_tokens: 0, tokenizer: TokenizerKind::default(), + strategy: ChunkStrategy::default(), } } } @@ -173,8 +216,12 @@ impl Default for ChunkConfig { /// Patterns always run BEFORE tokenization so that PII matches cannot be /// split across chunk boundaries. The scrubber emits an offset-delta map /// alongside the redacted text so chunk offsets can still be projected back -/// onto the original document for UI highlighting. -#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +/// onto the original document for UI highlighting, plus a [`PiiFinding`] +/// list recording every redaction with a confidence score and the contextual +/// anchors that fired. +/// +/// [`PiiFinding`]: crate::scrub::PiiFinding +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct ScrubProfile { /// Built-in pattern categories to enable. #[serde(default)] @@ -184,10 +231,57 @@ pub struct ScrubProfile { /// `RegexSet` pass alongside the built-ins. #[serde(default)] pub custom: Vec, + + /// Half-window size (in words) for the contextual anchor scan performed + /// around each candidate match. The scrubber looks `anchor_window` words + /// backward and forward for keyword anchors (e.g. `"ssn"`, `"credit card"`) + /// and boosts the confidence score accordingly. Default `7`. Set to `0` + /// to disable anchor boosting entirely. + #[serde(default = "default_anchor_window")] + pub anchor_window: u8, + + /// Minimum confidence `\[0.0, 1.0\]` for a candidate to be scrubbed and + /// reported. Candidates below this threshold are silently dropped + /// (neither replaced nor emitted into `pii_metadata`). Default `0.0` + /// (keep every candidate that passes algorithmic verification). + #[serde(default)] + pub min_confidence: f32, + + /// Entity slugs to **detect but not scrub** (report-only mode). The + /// finding is recorded in `pii_metadata` with its confidence and offsets, + /// but the original text passes through unchanged. Useful for audit-only + /// pipelines that need to know *where* PII is without redacting it. + /// + /// Example: `["email", "ssn"]` — emails and SSNs are detected and + /// reported, but the text is not masked. + #[serde(default)] + pub report_only: Vec, +} + +/// Default anchor half-window in words (the spec's recommended value). +#[must_use] +const fn default_anchor_window() -> u8 { + 7 +} + +impl Default for ScrubProfile { + fn default() -> Self { + Self { + patterns: Vec::new(), + custom: Vec::new(), + anchor_window: default_anchor_window(), + min_confidence: 0.0, + report_only: Vec::new(), + } + } } /// Built-in PII pattern categories. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +/// +/// `#[non_exhaustive]` allows new patterns to be added in minor releases +/// without breaking downstream exhaustive `match` arms. +#[non_exhaustive] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum BuiltInPattern { /// Email addresses (RFC-5322-ish). @@ -198,6 +292,8 @@ pub enum BuiltInPattern { Phone, /// Credit-card numbers: regex candidate followed by Luhn validation. CreditCard, + /// US bank ABA routing numbers (9 digits with checksum verification). + RoutingNumber, /// AWS access-key IDs (`AKIA...`) and secret access keys. AwsKey, /// GitHub personal access tokens (`ghp_...`, `gho_...`, etc.). @@ -269,10 +365,19 @@ impl SectionKind { /// `char_offset_start` / `char_offset_end` are offsets into the /// *post-scrubbed* document text. The scrubber's offset-delta map can /// project these back onto the original document for UI highlighting. -#[derive(Debug, Clone, PartialEq, Eq)] +/// +/// `pii` carries the [`PiiFinding`]s whose original-text ranges overlap this +/// chunk, for audit/logging in the `pii_metadata` Arrow column. +/// +/// [`PiiFinding`]: crate::scrub::PiiFinding +#[derive(Debug, Clone, PartialEq)] pub struct ChunkSpec { /// 0-based ordinal within the pipeline output. pub chunk_index: u32, + /// Deterministic content hash (`blake3` of text + source + offsets), + /// hex-encoded. Stable across identical pipeline runs for ID-based + /// dedup downstream. + pub chunk_id: String, /// The chunk text (post-scrubbing, post-tokenization slicing). pub text: String, /// BPE token count of `text`. Always `<= ChunkConfig::max_tokens`. @@ -285,7 +390,10 @@ pub struct ChunkSpec { pub section_kind: SectionKind, /// Half-open `[start, end)` character range into the scrubbed document. pub char_offset_start: u32, + /// Half-open `[start, end)` character range end into the scrubbed doc. pub char_offset_end: u32, + /// PII findings overlapping this chunk (offsets into original text). + pub pii: Vec, } #[cfg(test)] @@ -321,11 +429,15 @@ mod tests { regex: r"\bPROJECT-\d+\b".to_string(), replacement: "[PROJECT-ID]".to_string(), }], + ..ScrubProfile::default() }, chunk: ChunkConfig { max_tokens: 256, overlap_tokens: 32, tokenizer: TokenizerKind::O200kBase, + strategy: ChunkStrategy::Semantic { + similarity_threshold: 0.75, + }, }, source_label: Some("docs/architecture.md".to_string()), embeddings: None, @@ -379,10 +491,39 @@ mod tests { fn section_kind_all_covers_every_variant_without_dupes() { let all = SectionKind::all(); assert_eq!(all.len(), 7, "expected exactly 7 SectionKind variants"); - // Every variant must appear exactly once. for variant in all { let count = all.iter().filter(|&&v| v == *variant).count(); assert_eq!(count, 1, "duplicate variant in SectionKind::all()"); } } + + #[test] + fn chunk_strategy_round_trips_through_json() { + for original in [ + ChunkStrategy::Structural, + ChunkStrategy::Semantic { + similarity_threshold: 0.82, + }, + ] { + let json = serde_json::to_string(&original).expect("serialize"); + let back: ChunkStrategy = serde_json::from_str(&json).expect("deserialize"); + assert_eq!( + back, original, + "round-trip failed for {original:?} ({json})" + ); + } + // Wire shape: lowercase variant, snake_case field. + let s = serde_json::to_string(&ChunkStrategy::Semantic { + similarity_threshold: 0.5, + }) + .unwrap(); + assert_eq!(s, r#"{"semantic":{"similarity_threshold":0.5}}"#); + } + + #[test] + fn minimal_config_defaults_strategy_to_structural() { + let json = r#"{ "format": "markdown", "chunk": { "max_tokens": 128 } }"#; + let cfg: PipelineConfig = serde_json::from_str(json).unwrap(); + assert_eq!(cfg.chunk.strategy, ChunkStrategy::Structural); + } } diff --git a/crates/core/src/scrub.rs b/crates/core/src/scrub.rs index d38e28b..6626ec3 100644 --- a/crates/core/src/scrub.rs +++ b/crates/core/src/scrub.rs @@ -5,26 +5,52 @@ //! # Architecture //! //! 1. [`Scrubber::from_profile`] compiles all enabled patterns (built-in -//! and user-supplied) into individual [`Regex`](regex::Regex) instances. -//! 2. [`Scrubber::scrub`] finds all matches via `find_iter`, applies -//! optional post-filters (Luhn validation for credit cards), resolves -//! overlaps via greedy interval scheduling, and builds the scrubbed -//! text. +//! and user-supplied) into individual [`Regex`](regex::Regex) instances, +//! each tagged with its base confidence and contextual anchor keywords. +//! 2. [`Scrubber::scrub`] finds all matches via `find_iter`, pipes +//! credit-card candidates through Luhn mod-10 verification and routing +//! numbers through the ABA checksum (dropping failures), scans a ±N-word +//! window around each surviving candidate for keyword anchors, computes a +//! weighted-additive confidence score, resolves overlaps via greedy +//! interval scheduling, and builds the scrubbed text. //! 3. [`OffsetMap`] records every replacement's position so that chunk //! offsets (into the scrubbed text) can be projected back onto the //! original document for UI highlighting. +//! 4. [`PiiFinding`] records each redaction (entity, offset range, score, +//! anchors fired) into the `pii_metadata` Arrow column. +//! +//! # Confidence model +//! +//! Each candidate match receives a score in `[0.0, 1.0]`: +//! +//! `score = (base + anchors_hit × ANCHOR_DELTA).min(ANCHOR_CAP)` +//! +//! | Pattern | Base | Validator | Validator floor | +//! |-----------------|-------|--------------------|-----------------| +//! | Email | 0.85 | none | — | +//! | Ssn | 0.60 | none | — | +//! | Phone | 0.65 | none | — | +//! | `CreditCard` | 0.70 | Luhn (gate) | 0.90 on pass | +//! | `RoutingNumber` | 0.55 | ABA checksum (gate)| — | +//! | `AwsKey` | 0.95 | none | — | +//! | `GitHubPat` | 0.95 | none | — | +//! | `Jwt` | 0.90 | none | — | +//! +//! Candidates whose confidence falls below `ScrubProfile::min_confidence` +//! are dropped entirely (neither scrubbed nor reported). //! //! # Built-in patterns //! -//! | Pattern | Replacement | Post-filter | -//! |-------------|-----------------|-------------| -//! | Email | `[EMAIL]` | none | -//! | SSN | `[SSN]` | none | -//! | Phone | `[PHONE]` | none | -//! | CreditCard | `[CREDIT_CARD]` | Luhn | -//! | AwsKey | `[AWS_KEY]` | none | -//! | GitHubPat | `[GITHUB_PAT]` | none | -//! | Jwt | `[JWT]` | none | +//! | Pattern | Replacement | Post-filter | +//! |---------------|-------------------|-------------| +//! | Email | `[EMAIL]` | none | +//! | Ssn | `[SSN]` | none | +//! | Phone | `[PHONE]` | none | +//! | CreditCard | `[CREDIT_CARD]` | Luhn | +//! | RoutingNumber | `[ROUTING_NUMBER]`| ABA | +//! | AwsKey | `[AWS_KEY]` | none | +//! | GitHubPat | `[GITHUB_PAT]` | none | +//! | Jwt | `[JWT]` | none | use regex::Regex; @@ -32,6 +58,20 @@ use crate::error::{BitVanesError, Result}; use crate::parse::{Document, offset_to_u32}; use crate::schema::{BuiltInPattern, ScrubProfile}; +// =========================================================================== +// Confidence-model constants +// =========================================================================== + +/// Confidence added per contextual anchor hit. +const ANCHOR_DELTA: f32 = 0.10; + +/// Maximum confidence reachable via anchor boosting alone. +const ANCHOR_CAP: f32 = 0.99; + +/// Floor for Luhn-validated credit cards: a card that passes mod-10 is at +/// least this confident regardless of anchor context. +const LUHN_FLOOR: f32 = 0.90; + // =========================================================================== // Offset map // =========================================================================== @@ -138,6 +178,80 @@ impl OffsetMap { } } +// =========================================================================== +// PII findings (audit metadata) +// =========================================================================== + +/// A single PII finding emitted into the `pii_metadata` Arrow column. +/// +/// Offsets are byte offsets into the **original** (pre-scrub) document text, +/// so the audit UI can highlight the raw match directly. The `anchors_hit` +/// list records which contextual keywords fired within the surrounding word +/// window, making the confidence score explainable. +#[derive(Debug, Clone, PartialEq)] +pub struct PiiFinding { + /// Entity slug: `"email"`, `"ssn"`, `"credit_card"`, ... (or the lowercased + /// custom-pattern name). + pub entity: String, + /// Byte offset of the match start in the original document text. + pub offset_start: u32, + /// Byte offset of the match end (exclusive) in the original document text. + pub offset_end: u32, + /// Confidence score in `[0.0, 1.0]`. + pub confidence: f32, + /// Contextual anchor keywords that fired in the ±`anchor_window` word + /// window around this match. + pub anchors_hit: Vec, +} + +// =========================================================================== +// Anchor keyword tables +// =========================================================================== +// +// Lowercase. Single-word anchors are matched against whole-word tokens in +// the window; multi-word anchors (containing a space) are matched as +// substrings of the lowercased window. + +const ANCHOR_NONE: &[&str] = &[]; +const ANCHOR_EMAIL: &[&str] = &["email", "mail", "contact", "reach", "reply", "from", "to"]; +const ANCHOR_SSN: &[&str] = &[ + "ssn", + "social", + "security", + "tax", + "taxpayer", + "tin", + "identification", +]; +const ANCHOR_PHONE: &[&str] = &[ + "phone", + "mobile", + "cell", + "tel", + "telephone", + "fax", + "call", + "number", + "contact", +]; +const ANCHOR_CREDIT_CARD: &[&str] = &[ + "card", + "credit", + "debit", + "visa", + "mastercard", + "amex", + "discover", + "pan", + "cc", + "cvv", + "expiry", +]; +const ANCHOR_ROUTING: &[&str] = &["routing", "aba", "transit", "bank", "checking", "account"]; +const ANCHOR_AWS: &[&str] = &["aws", "access", "key", "iam", "credential", "secret"]; +const ANCHOR_GITHUB: &[&str] = &["github", "token", "pat", "gh", "gist"]; +const ANCHOR_JWT: &[&str] = &["jwt", "bearer", "token", "authorization", "auth"]; + // =========================================================================== // Scrubber // =========================================================================== @@ -149,6 +263,10 @@ impl OffsetMap { #[derive(Debug, Clone)] pub struct Scrubber { patterns: Vec, + anchor_window: u8, + min_confidence: f32, + /// Entity slugs to detect-and-report but NOT replace. + report_only: std::collections::HashSet, } #[derive(Debug, Clone)] @@ -156,27 +274,43 @@ struct CompiledPattern { regex: Regex, replacement: String, validator: Validator, + /// Stable slug emitted into [`PiiFinding::entity`]. + entity: &'static str, + /// Baseline confidence before anchor boosting. + base_confidence: f32, + /// Contextual anchor keywords for this pattern's matches. + anchors: &'static [&'static str], } -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] enum Validator { None, - /// Extract digits from the match, validate via Luhn checksum. + /// Extract digits from the match, validate via Luhn checksum. A failing + /// checksum drops the candidate entirely (not a credit card). Luhn, + /// Validate via the ABA routing-number checksum. A failing checksum drops + /// the candidate. + RoutingNumber, } impl Validator { + const fn is_gated(self) -> bool { + matches!(self, Self::Luhn | Self::RoutingNumber) + } + fn is_valid(self, text: &str, start: usize, end: usize) -> bool { match self { Self::None => true, Self::Luhn => luhn_valid(&text[start..end]), + Self::RoutingNumber => aba_valid(&text[start..end]), } } } impl Scrubber { /// Compiles all patterns (built-in and custom) from the profile into - /// a single [`Scrubber`]. + /// a single [`Scrubber`], inheriting the anchor-window size and + /// minimum-confidence threshold. /// /// # Errors /// @@ -185,7 +319,7 @@ impl Scrubber { let mut patterns = Vec::with_capacity(profile.patterns.len() + profile.custom.len()); for &kind in &profile.patterns { - let (src, replacement, validator) = builtin_config(kind); + let (src, replacement, validator, entity, base, anchors) = builtin_config(kind); let regex = Regex::new(src).map_err(|e| { BitVanesError::InvalidConfig(format!("built-in pattern {kind:?}: {e}")) })?; @@ -193,6 +327,9 @@ impl Scrubber { regex, replacement: replacement.to_string(), validator, + entity, + base_confidence: base, + anchors, }); } @@ -202,12 +339,20 @@ impl Scrubber { })?; patterns.push(CompiledPattern { regex, + entity: leak_entity(&custom.replacement), + base_confidence: 0.90, + anchors: ANCHOR_NONE, replacement: custom.replacement.clone(), validator: Validator::None, }); } - Ok(Self { patterns }) + Ok(Self { + patterns, + anchor_window: profile.anchor_window, + min_confidence: profile.min_confidence, + report_only: profile.report_only.iter().cloned().collect(), + }) } /// Returns `true` if no patterns are compiled (scrubbing is a no-op). @@ -216,37 +361,103 @@ impl Scrubber { self.patterns.is_empty() } - /// Scrubs `text`: finds all PII matches, resolves overlaps, and builds - /// the redacted output plus an [`OffsetMap`] for offset projection. + /// Scrubs `text`: finds all PII matches, runs algorithmic verification + /// (Luhn / ABA), scans the contextual anchor window, computes confidence + /// scores, resolves overlaps, and builds the redacted output. + /// + /// Returns `(scrubbed_text, offset_map, findings)`. Findings carry + /// offsets into the **original** text. #[must_use] - pub fn scrub(&self, text: &str) -> (String, OffsetMap) { + pub fn scrub(&self, text: &str) -> (String, OffsetMap, Vec) { if self.patterns.is_empty() || text.is_empty() { - return (text.to_string(), OffsetMap::default()); + return (text.to_string(), OffsetMap::default(), Vec::new()); } - let mut matches = Vec::new(); - for cp in &self.patterns { - for m in cp.regex.find_iter(text) { - if cp.validator.is_valid(text, m.start(), m.end()) { - matches.push(PiiMatch { - start: m.start(), - end: m.end(), - replacement: cp.replacement.clone(), - }); - } + let matches = self.find_matches(text); + let resolved = resolve_overlaps(matches); + build_scrubbed(text, &resolved) + } + + /// Runs all patterns against `text` and returns raw (pre-overlap-resolution) + /// matches. Parallelized behind the `parallel` feature. + fn find_matches(&self, text: &str) -> Vec { + #[cfg(feature = "parallel")] + { + if self.patterns.len() > 1 { + use rayon::prelude::*; + return self + .patterns + .par_iter() + .flat_map(|cp| self.find_pattern_matches(text, cp)) + .collect(); } } + // Sequential path (single pattern or no parallel feature). + self.patterns + .iter() + .flat_map(|cp| self.find_pattern_matches(text, cp)) + .collect() + } - let resolved = resolve_overlaps(matches); - build_scrubbed(text, &resolved) + /// Finds all matches for a single compiled pattern, with verification, + /// anchor scanning, and confidence scoring applied. + fn find_pattern_matches(&self, text: &str, cp: &CompiledPattern) -> Vec { + let mut out = Vec::new(); + for m in cp.regex.find_iter(text) { + let validator_ok = cp.validator.is_valid(text, m.start(), m.end()); + if cp.validator.is_gated() && !validator_ok { + continue; + } + let anchors_hit = + scan_anchors(text, m.start(), m.end(), self.anchor_window, cp.anchors); + let confidence = + compute_confidence(cp.base_confidence, anchors_hit.len(), cp.validator); + if confidence < self.min_confidence { + continue; + } + out.push(PiiMatch { + start: m.start(), + end: m.end(), + replacement: cp.replacement.clone(), + report_only: self.report_only.contains(cp.entity), + finding: PiiFinding { + entity: cp.entity.to_string(), + offset_start: offset_to_u32(m.start()), + offset_end: offset_to_u32(m.end()), + confidence, + anchors_hit: anchors_hit.iter().map(|s| (*s).to_string()).collect(), + }, + }); + } + out } } -/// A raw PII match found by a regex, before overlap resolution. +/// Derives a stable entity slug for a custom pattern from its replacement +/// token: strips bracket pairs and lowercases. Falls back to `"custom"`. +fn leak_entity(replacement: &str) -> &'static str { + // Leak is bounded: one tiny string per custom pattern per Scrubber build, + // and custom-pattern counts are small (user-supplied config). The + // alternative (changing `entity` to `String`) would force every built-in + // to allocate on every match — undesirable in the hot loop. + let trimmed = replacement.trim_matches(|c| c == '[' || c == ']'); + let slug = if trimmed.is_empty() { + "custom" + } else { + trimmed + }; + Box::leak(slug.to_ascii_lowercase().into_boxed_str()) +} + +/// A raw PII match found by a regex, before overlap resolution. Carries its +/// audit [`PiiFinding`] for downstream emission. struct PiiMatch { start: usize, end: usize, replacement: String, + /// If true, the finding is recorded but the original text is NOT replaced. + report_only: bool, + finding: PiiFinding, } /// Resolves overlapping matches using greedy interval scheduling (earliest @@ -271,10 +482,12 @@ fn resolve_overlaps(mut matches: Vec) -> Vec { } /// Interleaves original text with replacement tokens to produce the -/// scrubbed string and records each replacement in an [`OffsetMap`]. -fn build_scrubbed(text: &str, matches: &[PiiMatch]) -> (String, OffsetMap) { +/// scrubbed string, records each replacement in an [`OffsetMap`], and +/// collects the surviving [`PiiFinding`]s. +fn build_scrubbed(text: &str, matches: &[PiiMatch]) -> (String, OffsetMap, Vec) { let mut scrubbed = String::with_capacity(text.len()); let mut edits = Vec::with_capacity(matches.len()); + let mut findings = Vec::with_capacity(matches.len()); let mut cursor = 0usize; for m in matches { @@ -285,56 +498,236 @@ fn build_scrubbed(text: &str, matches: &[PiiMatch]) -> (String, OffsetMap) { if m.start > cursor { scrubbed.push_str(&text[cursor..m.start]); } - edits.push(OffsetEdit { - orig_start: m.start, - orig_end: m.end, - replacement_len: m.replacement.len(), - }); - scrubbed.push_str(&m.replacement); + if m.report_only { + // Detect-and-report: pass the original text through unchanged. + // No OffsetEdit (zero delta), but the finding is still recorded. + scrubbed.push_str(&text[m.start..m.end]); + } else { + edits.push(OffsetEdit { + orig_start: m.start, + orig_end: m.end, + replacement_len: m.replacement.len(), + }); + scrubbed.push_str(&m.replacement); + } cursor = m.end; + findings.push(m.finding.clone()); } if cursor < text.len() { scrubbed.push_str(&text[cursor..]); } - (scrubbed, OffsetMap { edits }) + (scrubbed, OffsetMap { edits }, findings) +} + +// =========================================================================== +// Contextual anchor scanning +// =========================================================================== + +/// Scans the ±`window_words` word window around `[start, end)` for keyword +/// anchors. Returns the (deduplicated, order-preserving) list of anchors found. +/// +/// Single-word anchors require a whole-word match; multi-word anchors match +/// as substrings of the lowercased window. ASCII-whitespace byte boundaries +/// are always UTF-8 char boundaries, so the window slice is always valid. +#[must_use] +fn scan_anchors( + text: &str, + start: usize, + end: usize, + window_words: u8, + anchors: &'static [&'static str], +) -> Vec<&'static str> { + if anchors.is_empty() || window_words == 0 { + return Vec::new(); + } + let words = usize::from(window_words); + let win_start = window_back(text, start, words); + let win_end = window_forward(text, end, words); + if win_start >= win_end { + return Vec::new(); + } + let window = &text[win_start..win_end]; + let lower = window.to_ascii_lowercase(); + let word_set: std::collections::HashSet<&str> = + lower.split(|c: char| !c.is_alphanumeric()).collect(); + + let mut hits: Vec<&'static str> = Vec::new(); + for &anchor in anchors { + if anchor.contains(' ') { + if lower.contains(anchor) && !hits.contains(&anchor) { + hits.push(anchor); + } + } else if word_set.contains(anchor) && !hits.contains(&anchor) { + hits.push(anchor); + } + } + hits +} + +/// Walks backward from `from`, returning the byte offset that begins the +/// `words`-th word before `from`. Clamps to `0`. +fn window_back(text: &str, from: usize, words: usize) -> usize { + let bytes = text.as_bytes(); + if words == 0 || from == 0 { + return from.min(bytes.len()); + } + let mut i = from.min(bytes.len()); + let mut count = 0usize; + // Skip whitespace we're currently sitting in. + while i > 0 && bytes[i - 1].is_ascii_whitespace() { + i -= 1; + } + while count < words { + // Consume one word (non-whitespace run). + while i > 0 && !bytes[i - 1].is_ascii_whitespace() { + i -= 1; + } + count += 1; + if count >= words { + break; + } + // Skip preceding whitespace before the next word. + while i > 0 && bytes[i - 1].is_ascii_whitespace() { + i -= 1; + } + } + i +} + +/// Walks forward from `from`, returning the byte offset that ends the +/// `words`-th word after `from`. Clamps to `text.len()`. +fn window_forward(text: &str, from: usize, words: usize) -> usize { + let bytes = text.as_bytes(); + let n = bytes.len(); + if words == 0 { + return from.min(n); + } + let mut i = from.min(n); + let mut count = 0usize; + // Skip whitespace we're currently sitting in. + while i < n && bytes[i].is_ascii_whitespace() { + i += 1; + } + while i < n && count < words { + // Consume one word. + while i < n && !bytes[i].is_ascii_whitespace() { + i += 1; + } + count += 1; + // Consume trailing whitespace. + while i < n && bytes[i].is_ascii_whitespace() { + i += 1; + } + } + i +} + +// =========================================================================== +// Confidence scoring +// =========================================================================== + +/// Weighted-additive confidence: `base + anchors × delta`, gated by +/// validator. For Luhn-validated cards, a passing checksum floors the score +/// at [`LUHN_FLOOR`]; a failing checksum is unreachable here (gated out +/// upstream) and returns `0.0` defensively. +#[must_use] +fn compute_confidence(base: f32, anchors_hit: usize, validator: Validator) -> f32 { + let boost = f32::from(u8::try_from(anchors_hit.min(255)).unwrap_or(255)) * ANCHOR_DELTA; + match validator { + Validator::Luhn => base.max(LUHN_FLOOR) + boost, + Validator::RoutingNumber | Validator::None => (base + boost).min(ANCHOR_CAP), + } + .min(ANCHOR_CAP) } // =========================================================================== // Built-in pattern definitions // =========================================================================== -/// Returns `(regex_source, replacement_token, validator)` for a built-in kind. -fn builtin_config(kind: BuiltInPattern) -> (&'static str, &'static str, Validator) { +/// Returns `(regex_source, replacement_token, validator, entity_slug, +/// base_confidence, anchor_keywords)` for a built-in kind. +#[allow(clippy::type_complexity)] +fn builtin_config( + kind: BuiltInPattern, +) -> ( + &'static str, + &'static str, + Validator, + &'static str, + f32, + &'static [&'static str], +) { match kind { BuiltInPattern::Email => ( r"[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}", "[EMAIL]", Validator::None, + "email", + 0.85, + ANCHOR_EMAIL, + ), + BuiltInPattern::Ssn => ( + r"\b\d{3}-\d{2}-\d{4}\b", + "[SSN]", + Validator::None, + "ssn", + 0.60, + ANCHOR_SSN, + ), + BuiltInPattern::Phone => ( + r"\+1\d{10}\b", + "[PHONE]", + Validator::None, + "phone", + 0.65, + ANCHOR_PHONE, ), - BuiltInPattern::Ssn => (r"\b\d{3}-\d{2}-\d{4}\b", "[SSN]", Validator::None), - BuiltInPattern::Phone => (r"\+1\d{10}\b", "[PHONE]", Validator::None), BuiltInPattern::CreditCard => ( r"\b(?:\d[ -]?){12,18}\d\b", "[CREDIT_CARD]", Validator::Luhn, + "credit_card", + 0.70, + ANCHOR_CREDIT_CARD, + ), + BuiltInPattern::RoutingNumber => ( + r"\b\d{9}\b", + "[ROUTING_NUMBER]", + Validator::RoutingNumber, + "routing_number", + 0.55, + ANCHOR_ROUTING, + ), + BuiltInPattern::AwsKey => ( + r"\bAKIA[0-9A-Z]{16}\b", + "[AWS_KEY]", + Validator::None, + "aws_key", + 0.95, + ANCHOR_AWS, ), - BuiltInPattern::AwsKey => (r"\bAKIA[0-9A-Z]{16}\b", "[AWS_KEY]", Validator::None), BuiltInPattern::GitHubPat => ( r"\bgh[pousr]_[A-Za-z0-9]{36,}\b", "[GITHUB_PAT]", Validator::None, + "github_pat", + 0.95, + ANCHOR_GITHUB, ), BuiltInPattern::Jwt => ( r"\beyJ[A-Za-z0-9_\-]*\.eyJ[A-Za-z0-9_\-]*\.[A-Za-z0-9_\-]*\b", "[JWT]", Validator::None, + "jwt", + 0.90, + ANCHOR_JWT, ), } } // =========================================================================== -// Luhn checksum +// Checksums: Luhn (credit cards) + ABA (routing numbers) // =========================================================================== /// Validates a string of digits (possibly mixed with separators) using the @@ -364,6 +757,25 @@ fn luhn_valid(text: &str) -> bool { sum % 10 == 0 } +/// Validates a US ABA routing number (9 digits) using the standard +/// fractional-checksum: `3·(d0+d3+d6) + 7·(d1+d4+d7) + (d2+d5+d8)` must be +/// divisible by 10. Returns `false` if the extracted digit count is not 9. +fn aba_valid(text: &str) -> bool { + let digits: Vec = text + .bytes() + .filter(u8::is_ascii_digit) + .map(|b| b - b'0') + .collect(); + if digits.len() != 9 { + return false; + } + let d = &digits[..9]; + let sum = 3 * u32::from(d[0] + d[3] + d[6]) + + 7 * u32::from(d[1] + d[4] + d[7]) + + u32::from(d[2] + d[5] + d[8]); + sum % 10 == 0 +} + // =========================================================================== // High-level entry points // =========================================================================== @@ -375,7 +787,10 @@ fn luhn_valid(text: &str) -> bool { /// /// Returns [`BitVanesError::InvalidConfig`] if any regex in the profile /// fails to compile. -pub fn scrub_text(text: &str, profile: &ScrubProfile) -> Result<(String, OffsetMap)> { +pub fn scrub_text( + text: &str, + profile: &ScrubProfile, +) -> Result<(String, OffsetMap, Vec)> { let scrubber = Scrubber::from_profile(profile)?; Ok(scrubber.scrub(text)) } @@ -384,14 +799,18 @@ pub fn scrub_text(text: &str, profile: &ScrubProfile) -> Result<(String, OffsetM /// the scrubbed text's coordinate space. /// /// This is the primary entry point called by the pipeline between parsing -/// and chunking. +/// and chunking. Returns the scrubbed document, the offset map, and the +/// findings (offsets into the original text). /// /// # Errors /// /// Returns [`BitVanesError::InvalidConfig`] if any regex fails to compile. -pub fn scrub_document(doc: Document, profile: &ScrubProfile) -> Result<(Document, OffsetMap)> { +pub fn scrub_document( + doc: Document, + profile: &ScrubProfile, +) -> Result<(Document, OffsetMap, Vec)> { let scrubber = Scrubber::from_profile(profile)?; - let (scrubbed_text, map) = scrubber.scrub(&doc.full_text); + let (scrubbed_text, map, findings) = scrubber.scrub(&doc.full_text); let spans = doc .spans @@ -411,6 +830,7 @@ pub fn scrub_document(doc: Document, profile: &ScrubProfile) -> Result<(Document spans, }, map, + findings, )) } @@ -419,11 +839,13 @@ mod tests { use super::*; use crate::parse::Parser; use crate::schema::{CustomPattern, DocumentFormat, PipelineConfig}; + use proptest::prop_assert; - fn scrub_with(text: &str, patterns: &[BuiltInPattern]) -> (String, OffsetMap) { + fn scrub_with(text: &str, patterns: &[BuiltInPattern]) -> (String, OffsetMap, Vec) { let scrubber = Scrubber::from_profile(&ScrubProfile { patterns: patterns.to_vec(), custom: vec![], + ..ScrubProfile::default() }) .expect("built-in patterns should compile"); scrubber.scrub(text) @@ -433,53 +855,58 @@ mod tests { #[test] fn email_is_redacted() { - let (out, map) = scrub_with( + let (out, map, findings) = scrub_with( "Contact alice@example.com for details.", &[BuiltInPattern::Email], ); assert_eq!(out, "Contact [EMAIL] for details."); assert_eq!(map.len(), 1); + assert_eq!(findings.len(), 1); + assert_eq!(findings[0].entity, "email"); + assert!(findings[0].confidence >= 0.85, "email base confidence"); } #[test] fn multiple_emails_in_one_text() { - let (out, map) = scrub_with( + let (out, map, findings) = scrub_with( "From a@x.com and b@y.org to c@z.io.", &[BuiltInPattern::Email], ); assert_eq!(out, "From [EMAIL] and [EMAIL] to [EMAIL]."); assert_eq!(map.len(), 3); + assert_eq!(findings.len(), 3); } #[test] fn ssn_is_redacted() { - let (out, _) = scrub_with("SSN: 123-45-6789.", &[BuiltInPattern::Ssn]); + let (out, _, findings) = scrub_with("SSN: 123-45-6789.", &[BuiltInPattern::Ssn]); assert_eq!(out, "SSN: [SSN]."); + assert_eq!(findings[0].entity, "ssn"); } #[test] fn phone_e164_is_redacted() { - let (out, _) = scrub_with("Call +15551234567.", &[BuiltInPattern::Phone]); + let (out, _, _) = scrub_with("Call +15551234567.", &[BuiltInPattern::Phone]); assert_eq!(out, "Call [PHONE]."); } #[test] fn aws_key_is_redacted() { - let (out, _) = scrub_with("key = AKIAIOSFODNN7EXAMPLE", &[BuiltInPattern::AwsKey]); + let (out, _, _) = scrub_with("key = AKIAIOSFODNN7EXAMPLE", &[BuiltInPattern::AwsKey]); assert_eq!(out, "key = [AWS_KEY]"); } #[test] fn github_pat_is_redacted() { let pat = "ghp_abcdefghijklmnopqrstuvwxyz0123456789AB"; - let (out, _) = scrub_with(&format!("token: {pat}"), &[BuiltInPattern::GitHubPat]); + let (out, _, _) = scrub_with(&format!("token: {pat}"), &[BuiltInPattern::GitHubPat]); assert_eq!(out, "token: [GITHUB_PAT]"); } #[test] fn jwt_is_redacted() { let jwt = "eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c"; - let (out, _) = scrub_with( + let (out, _, _) = scrub_with( &format!("Authorization: Bearer {jwt}"), &[BuiltInPattern::Jwt], ); @@ -495,18 +922,31 @@ mod tests { #[test] fn valid_credit_card_passes_luhn() { // 4111 1111 1111 1111 is a classic test card (passes Luhn). - let (out, map) = scrub_with( + let (out, map, findings) = scrub_with( "Card: 4111 1111 1111 1111 done.", &[BuiltInPattern::CreditCard], ); assert_eq!(out, "Card: [CREDIT_CARD] done."); assert_eq!(map.len(), 1); + assert_eq!(findings[0].entity, "credit_card"); + // Luhn-passing cards are floored at LUHN_FLOOR (0.90). + assert!( + findings[0].confidence >= 0.90, + "Luhn-valid card confidence {} should be >= 0.90", + findings[0].confidence + ); + // "Card" anchor should fire in the window. + assert!( + findings[0].anchors_hit.iter().any(|a| a == "card"), + "expected 'card' anchor to fire, got {:?}", + findings[0].anchors_hit + ); } #[test] fn invalid_credit_card_fails_luhn_and_is_not_redacted() { // Same length but invalid checksum. - let (out, map) = scrub_with( + let (out, map, findings) = scrub_with( "Card: 4111 1111 1111 1112 done.", &[BuiltInPattern::CreditCard], ); @@ -515,21 +955,171 @@ mod tests { "invalid-Luhn sequence should NOT be redacted" ); assert!(map.is_empty()); + assert!( + findings.is_empty(), + "no finding should be emitted for a failed gate" + ); } #[test] fn credit_card_without_separators() { // 4111111111111111 — unseparated, valid Luhn. - let (out, _) = scrub_with("4111111111111111", &[BuiltInPattern::CreditCard]); + let (out, _, _) = scrub_with("4111111111111111", &[BuiltInPattern::CreditCard]); assert_eq!(out, "[CREDIT_CARD]"); } + // ----- Routing number + ABA ----- + + #[test] + fn valid_routing_number_passes_aba() { + // 021000021 is a real Bank of America ABA routing number. + let (out, map, findings) = scrub_with( + "Bank routing 021000021 for checks.", + &[BuiltInPattern::RoutingNumber], + ); + assert_eq!(out, "Bank routing [ROUTING_NUMBER] for checks."); + assert_eq!(map.len(), 1); + assert_eq!(findings[0].entity, "routing_number"); + // "routing" and "bank" anchors should fire. + assert!(!findings[0].anchors_hit.is_empty()); + } + + #[test] + fn invalid_routing_number_fails_aba() { + // 9 digits but invalid ABA checksum (last digit wrong). + let (out, map, findings) = + scrub_with("Number 123456789 here.", &[BuiltInPattern::RoutingNumber]); + assert_eq!( + out, "Number 123456789 here.", + "invalid ABA should NOT be redacted" + ); + assert!(map.is_empty()); + assert!(findings.is_empty()); + } + + // ----- Confidence & anchor scoring ----- + + #[test] + fn anchor_boost_raises_confidence() { + // Same SSN, but one has contextual anchors and the other doesn't. + let (_, _, findings_with) = scrub_with( + "Social security number: 123-45-6789 is the SSN.", + &[BuiltInPattern::Ssn], + ); + let (_, _, findings_without) = scrub_with( + "Random token 123-45-6789 appears here.", + &[BuiltInPattern::Ssn], + ); + assert_eq!(findings_with.len(), 1); + assert_eq!(findings_without.len(), 1); + assert!( + findings_with[0].confidence > findings_without[0].confidence, + "anchored match ({}) should outrank unanchored ({})", + findings_with[0].confidence, + findings_without[0].confidence + ); + assert!(!findings_with[0].anchors_hit.is_empty()); + } + + #[test] + fn confidence_respects_cap() { + // Many anchors should still cap at ANCHOR_CAP. + let (_, _, findings) = scrub_with( + "ssn social security tax taxpayer tin identification number: 123-45-6789", + &[BuiltInPattern::Ssn], + ); + assert_eq!(findings.len(), 1); + assert!( + findings[0].confidence <= ANCHOR_CAP + 1e-6, + "confidence {} must not exceed ANCHOR_CAP", + findings[0].confidence + ); + } + + #[test] + fn zero_anchor_window_disables_boosting() { + let scrubber = Scrubber::from_profile(&ScrubProfile { + patterns: vec![BuiltInPattern::Ssn], + custom: vec![], + anchor_window: 0, + min_confidence: 0.0, + ..ScrubProfile::default() + }) + .unwrap(); + let (_, _, findings) = scrubber.scrub("SSN social security 123-45-6789"); + assert_eq!(findings.len(), 1); + assert!( + findings[0].anchors_hit.is_empty(), + "anchor_window=0 must produce no anchor hits" + ); + // Base confidence for SSN is 0.60, no boost. + assert!((findings[0].confidence - 0.60).abs() < 1e-6); + } + + #[test] + fn min_confidence_filters_low_score_matches() { + // SSN without anchors has confidence 0.60. Set threshold to 0.65. + let scrubber = Scrubber::from_profile(&ScrubProfile { + patterns: vec![BuiltInPattern::Ssn], + custom: vec![], + anchor_window: 7, + min_confidence: 0.65, + ..ScrubProfile::default() + }) + .unwrap(); + let (out, map, findings) = scrubber.scrub("Random token 123-45-6789 appears."); + assert!( + !out.contains("[SSN]"), + "below-threshold match should not be scrubbed: {out}" + ); + assert!(map.is_empty()); + assert!(findings.is_empty()); + + // With anchors the confidence climbs above 0.65 and the match survives. + let (out2, _, findings2) = scrubber.scrub("Social security number 123-45-6789 is here."); + assert!( + out2.contains("[SSN]"), + "above-threshold match should be scrubbed" + ); + assert_eq!(findings2.len(), 1); + } + + #[test] + fn report_only_detects_but_does_not_scrub() { + let scrubber = Scrubber::from_profile(&ScrubProfile { + patterns: vec![BuiltInPattern::Email, BuiltInPattern::Ssn], + custom: vec![], + report_only: vec!["email".to_string()], + ..ScrubProfile::default() + }) + .unwrap(); + + let (out, map, findings) = scrubber.scrub("Email alice@test.com, SSN 123-45-6789."); + // Email should NOT be scrubbed (report-only), SSN SHOULD be scrubbed. + assert!( + out.contains("alice@test.com"), + "report-only email should pass through: {out}" + ); + assert!(out.contains("[SSN]"), "SSN should still be scrubbed: {out}"); + assert!(!out.contains("123-45-6789"), "SSN should be redacted"); + // Both findings should be recorded. + assert_eq!(findings.len(), 2, "both findings should be in metadata"); + assert_eq!(findings[0].entity, "email"); + assert_eq!(findings[1].entity, "ssn"); + // OffsetMap should only have one edit (SSN, not email). + assert_eq!( + map.len(), + 1, + "only the scrubbed match creates an offset edit" + ); + } + // ----- OffsetMap projection ----- #[test] fn offset_map_forward_round_trip() { let text = "Contact alice@example.com please."; - let (scrubbed, map) = scrub_with(text, &[BuiltInPattern::Email]); + let (scrubbed, map, _) = scrub_with(text, &[BuiltInPattern::Email]); assert_eq!(scrubbed, "Contact [EMAIL] please."); // [EMAIL] occupies scrubbed positions [8..15). Positions OUTSIDE the @@ -546,7 +1136,7 @@ mod tests { #[test] fn offset_map_identity_when_no_matches() { - let (out, map) = scrub_with("No PII here.", &[BuiltInPattern::Email]); + let (out, map, _) = scrub_with("No PII here.", &[BuiltInPattern::Email]); assert!(map.is_identity()); assert_eq!(out, "No PII here."); // Forward projection of any offset is the offset itself. @@ -556,6 +1146,20 @@ mod tests { } } + #[test] + fn findings_offsets_are_into_original_text() { + let original = "Email alice@example.com now."; + let (out, _, findings) = scrub_with(original, &[BuiltInPattern::Email]); + assert_eq!(out, "Email [EMAIL] now."); + assert_eq!(findings.len(), 1); + let f = &findings[0]; + let matched = &original[f.offset_start as usize..f.offset_end as usize]; + assert_eq!( + matched, "alice@example.com", + "finding offsets must reference the original text" + ); + } + // ----- scrub_document ----- #[test] @@ -572,11 +1176,12 @@ mod tests { let profile = ScrubProfile { patterns: vec![BuiltInPattern::Email], - custom: vec![], + ..ScrubProfile::default() }; - let (scrubbed_doc, map) = scrub_document(doc.clone(), &profile).expect("scrub"); + let (scrubbed_doc, map, findings) = scrub_document(doc.clone(), &profile).expect("scrub"); assert_eq!(map.len(), 2, "two emails should be scrubbed"); + assert_eq!(findings.len(), 2, "two findings should be emitted"); assert!(scrubbed_doc.full_text.contains("[EMAIL]")); assert!(!scrubbed_doc.full_text.contains("alice@test.com")); @@ -606,9 +1211,10 @@ mod tests { .parse("Just some text.\n\nNo PII.", &cfg) .expect("parse"); - let (scrubbed_doc, map) = + let (scrubbed_doc, map, findings) = scrub_document(doc.clone(), &ScrubProfile::default()).expect("scrub"); assert!(map.is_identity()); + assert!(findings.is_empty()); assert_eq!(scrubbed_doc, doc); } @@ -622,11 +1228,19 @@ mod tests { regex: r"PROJECT-\d{4}".to_string(), replacement: "[PROJ]".to_string(), }], + ..ScrubProfile::default() }) .expect("custom regex should compile"); - let (out, _) = scrubber.scrub("See PROJECT-1234 and PROJECT-5678."); + let (out, _, findings) = scrubber.scrub("See PROJECT-1234 and PROJECT-5678."); assert_eq!(out, "See [PROJ] and [PROJ]."); + assert_eq!(findings.len(), 2); + // Entity derived from replacement token "[PROJ]" → "proj". + assert_eq!(findings[0].entity, "proj"); + assert!( + findings[0].confidence >= 0.90, + "custom patterns default to 0.90 base" + ); } #[test] @@ -637,6 +1251,7 @@ mod tests { regex: "[invalid".to_string(), replacement: "X".to_string(), }], + ..ScrubProfile::default() }) .unwrap_err(); assert!( @@ -656,11 +1271,12 @@ mod tests { regex: r"alice@example\.com\.[a-z]+".to_string(), replacement: "[ALICE_FULL]".to_string(), }], + ..ScrubProfile::default() }) .expect("compile"); // The custom pattern is longer and overlaps the email match. - let (out, map) = scrubber.scrub("Contact alice@example.com.org now."); + let (out, map, findings) = scrubber.scrub("Contact alice@example.com.org now."); // Greedy scheduling: earliest-ending wins. Email ends at // "alice@example.com" which is earlier than the custom match // "alice@example.com.org". @@ -669,9 +1285,10 @@ mod tests { "one pattern should win: {out}" ); assert_eq!(map.len(), 1, "exactly one match should survive overlap"); + assert_eq!(findings.len(), 1); } - // ----- Luhn unit tests ----- + // ----- Checksum unit tests ----- #[test] fn luhn_known_valid_cards() { @@ -693,4 +1310,110 @@ mod tests { assert!(!luhn_valid("12345")); assert!(!luhn_valid("")); } + + #[test] + fn aba_known_valid_routing_numbers() { + assert!(aba_valid("021000021")); // Bank of America + assert!(aba_valid("026009593")); // Bank of America (different branch) + assert!(aba_valid("111000025")); // Wells Fargo + } + + #[test] + fn aba_rejects_invalid_checksums() { + assert!(!aba_valid("123456789")); + assert!(!aba_valid("999999999")); + assert!(!aba_valid("021000022")); // last digit wrong + } + + #[test] + fn aba_rejects_wrong_length() { + assert!(!aba_valid("12345678")); + assert!(!aba_valid("1234567890")); + assert!(!aba_valid("")); + } + + // ----- Anchor window helpers ----- + + #[test] + fn window_back_handles_whitespace() { + let text = "foo bar baz qux"; + // "qux" starts at byte 12. Going back 1 word → 8 ("baz"). + assert_eq!(window_back(text, 12, 1), 8); + // Going back 2 words → 4 ("bar"). + assert_eq!(window_back(text, 12, 2), 4); + // Going back more words than exist → clamps to 0. + assert_eq!(window_back(text, 12, 99), 0); + } + + #[test] + fn window_forward_handles_whitespace() { + let text = "foo bar baz qux"; + // From byte 0 ("foo"), forward 1 word → 4 (after "foo "). + assert_eq!(window_forward(text, 0, 1), 4); + // Forward 2 words → 8. + assert_eq!(window_forward(text, 0, 2), 8); + // Forward beyond text → clamps to len. + assert_eq!(window_forward(text, 0, 99), text.len()); + } + + // ----- Property tests (proptest) ----- + + proptest::proptest! { + /// OffsetMap invariant: for any scrubbed-text offset OUTSIDE a + /// replaced region, forward(inverse(s)) == s. We can't guarantee this + /// inside a replacement token (the map intentionally snaps), so we + /// skip those ranges by checking that the round-trip produces a + /// valid offset in the same "gap" of the scrubbed text. + #[test] + fn offsetmap_roundtrip_preserves_gaps( + words in proptest::collection::vec(r"[a-z]{1,8}", 1..40) + ) { + // Insert fake "emails" so the scrubber has something to replace. + let text = words.iter() + .flat_map(|w| [w.as_str(), " ", "x@y.z", " "]) + .collect::(); + let (scrubbed, map, _) = scrub_with(&text, &[BuiltInPattern::Email]); + + // Every position in the scrubbed text should project to a valid + // original-text offset and back without panicking or going + // out-of-bounds. + for s in 0..=scrubbed.len() { + let orig = map.project_inverse(s); + prop_assert!(orig <= text.len(), "inverse({s}) = {orig} > text.len()"); + let fwd = map.project_forward(orig); + prop_assert!(fwd <= scrubbed.len(), "forward({orig}) = {fwd} > scrubbed.len()"); + } + } + + /// Every emitted finding has: valid offset bounds, confidence in + /// [0, 1], and offset_end > offset_start. + #[test] + fn findings_satisfy_bounds_invariant( + words in proptest::collection::vec(r"[a-zA-Z0-9@._%+\-]{1,20}", 1..60) + ) { + let text = words.join(" "); + let profile = ScrubProfile { + patterns: vec![ + BuiltInPattern::Email, + BuiltInPattern::CreditCard, + BuiltInPattern::Ssn, + ], + ..ScrubProfile::default() + }; + let scrubber = Scrubber::from_profile(&profile).unwrap(); + let (scrubbed, map, findings) = scrubber.scrub(&text); + + for f in &findings { + prop_assert!(f.offset_end > f.offset_start, "empty finding range"); + prop_assert!(f.offset_end as usize <= text.len(), "offset past text end"); + prop_assert!(f.confidence >= 0.0 && f.confidence <= 1.0, "confidence out of [0,1]"); + } + // The scrubbed text must never be longer than the original + // (replacements are tokens like [EMAIL] which are ≤ the matched PII + // in practice, but the real invariant is: scrubbed.len() and + // map agree on the final offset). + let _ = scrubbed; + let _ = map; + } + } } diff --git a/crates/wasm/Cargo.toml b/crates/wasm/Cargo.toml index 3f2ceb5..5fabf62 100644 --- a/crates/wasm/Cargo.toml +++ b/crates/wasm/Cargo.toml @@ -15,7 +15,7 @@ description = "Thin wasm-bindgen wrapper exposing bitvanes-core to JavaScript vi crate-type = ["cdylib", "rlib"] [dependencies] -bitvanes-core = { path = "../core", version = "0.1.1" } +bitvanes-core = { path = "../core", version = "0.4.0" } console_error_panic_hook = "0.1.7" serde = { version = "1.0.228", features = ["derive"] } serde-wasm-bindgen = "0.6.5" diff --git a/crates/wasm/src/lib.rs b/crates/wasm/src/lib.rs index 31904b9..8c655ca 100644 --- a/crates/wasm/src/lib.rs +++ b/crates/wasm/src/lib.rs @@ -107,7 +107,9 @@ pub fn active_export_count() -> u32 { /// Processes a document and returns chunk data as a JS array (via serde, /// NOT zero-copy Arrow FFI). Fallback for when `parseRecordBatch` fails. /// -/// Each element is `{ chunk_index, text, token_count, heading_path, section_kind }`. +/// Each element carries the chunk text, token count, heading ancestry, +/// section kind, deterministic `chunk_id`, and the PII findings whose +/// original-text ranges overlap the chunk. #[wasm_bindgen] pub fn process_chunks(config_js: JsValue, bytes: &[u8]) -> Result { let cfg: bitvanes_core::PipelineConfig = serde_wasm_bindgen::from_value(config_js) @@ -115,20 +117,33 @@ pub fn process_chunks(config_js: JsValue, bytes: &[u8]) -> Result = chunks .iter() .map(|c| SimpleChunk { chunk_index: c.chunk_index, + chunk_id: c.chunk_id.clone(), text: c.text.clone(), token_count: c.token_count, heading_path: c.heading_path.clone(), section_kind: format!("{:?}", c.section_kind).to_lowercase(), + pii: c + .pii + .iter() + .map(|f| SimpleFinding { + entity: f.entity.clone(), + offset_start: f.offset_start, + offset_end: f.offset_end, + confidence: f.confidence, + anchors_hit: f.anchors_hit.clone(), + }) + .collect(), }) .collect(); @@ -138,8 +153,19 @@ pub fn process_chunks(config_js: JsValue, bytes: &[u8]) -> Result, section_kind: String, + pii: Vec, +} + +#[derive(serde::Serialize)] +struct SimpleFinding { + entity: String, + offset_start: u32, + offset_end: u32, + confidence: f32, + anchors_hit: Vec, }