From ac1a8db7ac8e5a5f29a817d95972aa7833d95937 Mon Sep 17 00:00:00 2001 From: Santiago Rivera <87353936+1816x@users.noreply.github.com> Date: Tue, 11 Aug 2026 12:47:47 -0600 Subject: [PATCH] Add durable HNSW and corpus snapshots --- README.md | 39 +++- app/app/lib/rag.ts | 1 + bindings/python/hnsw_rag/embeddings.py | 1 + bindings/src/lib.rs | 34 ++- bindings/tests/test_hnsw.py | 17 ++ engine/src/hnsw.rs | 287 ++++++++++++++++++++++++- engine/src/lib.rs | 2 +- engine/src/rng.rs | 4 + engine/tests/snapshot.rs | 83 +++++++ service/rag_service/app.py | 7 +- service/rag_service/store.py | 282 ++++++++++++++++++------ service/tests/test_e2e.py | 5 + service/tests/test_persistence.py | 51 +++++ 13 files changed, 742 insertions(+), 71 deletions(-) create mode 100644 engine/tests/snapshot.rs create mode 100644 service/tests/test_persistence.py diff --git a/README.md b/README.md index dfe04da..26d5241 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,7 @@ fly launch --no-deploy # claims an app name, keeps the committed fly.toml fly deploy ``` -The `Dockerfile` is multi-stage — stage one compiles the PyO3 wheel with the Rust toolchain, stage two installs only that wheel (it's `abi3`, so it's portable across CPython ≥ 3.9). Final image is ~270 MB and runs as a non-root user. `fly.toml` scales to zero when idle and suspends rather than stops, so the in-memory index survives a resume. +The `Dockerfile` is multi-stage — stage one compiles the PyO3 wheel with the Rust toolchain, stage two installs only that wheel (it's `abi3`, so it's portable across CPython ≥ 3.9). Final image is ~270 MB and runs as a non-root user. `fly.toml` scales to zero when idle and suspends rather than stops, so the default in-memory index survives a resume. Durable snapshots can be enabled explicitly as described below. **Frontend (Vercel):** deploy `app/` as its own project. Point it at the backend with `RAG_SERVICE_URL`, or set `DEFAULT_RAG_SERVICE_URL` in `app/app/lib/config.ts` (it's a public URL, not a secret). @@ -47,7 +47,7 @@ The `Dockerfile` is multi-stage — stage one compiles the PyO3 wheel with the R - **No `ANTHROPIC_API_KEY`.** It runs in mock mode: answers are extractive, tagged with a visible `mock` badge. Nothing calls Claude, so there's no key on a public endpoint and no spend to burn. Retrieval — the HNSW index, which is the point of the project — is fully real. - **The hashed fallback embedder**, not `sentence-transformers` (which would drag `torch` into the image). That means retrieval matches on *term overlap, not meaning*. Don't mistake the demo for semantic search; install `sentence-transformers` and set `RAG_EMBEDDER=model` for that. `GET /stats` reports which backend is live so you never have to guess. - **Cold starts.** Scaled to zero, the first request after an idle period waits a second or two for the machine to wake. -- **A visible, locked document workspace.** The UI lists every indexed document and shows the active corpus, embedding, HNSW, relevance, generation, and upload configuration. Public uploads are disabled by default; enabling them is an explicit deployment choice. Documents are held in memory and disappear when the service starts fresh. +- **A visible, locked document workspace.** The UI lists every indexed document and shows the active corpus, embedding, HNSW, relevance, generation, and upload configuration. Public uploads are disabled by default; enabling them is an explicit deployment choice. The hosted configuration remains in-memory; deployments can opt into durable local snapshots with `RAG_STATE_PATH`. ## Quick start @@ -75,6 +75,8 @@ from hnsw_engine import Hnsw idx = Hnsw(dim=384, metric="cosine") idx.insert_batch(vectors) # -> [0, 1, 2, ...] idx.search(query, k=10) # -> [(id, distance), ...] closest first +idx.save("index.hnsw") # atomic, versioned binary snapshot +idx = Hnsw.load("index.hnsw") # checked restore; inserts may continue ``` **3. Full RAG (service + app):** the service runs **without an API key** in mock mode, so you can see the whole pipeline before adding one. @@ -114,10 +116,43 @@ on disk. `scripts/seed.py` uses the same protected upload endpoint and therefore also requires uploads to be enabled. Startup seeding is unaffected because it inserts directly into the in-process store. +### Optional durable persistence + +The default remains process-local and in-memory. Set `RAG_STATE_PATH` to a writable +file (for example `/data/corpus.rag`) to opt in. If the file does not exist the +service starts empty and creates it after the first successful insertion. If it +does exist, startup restores the HNSW vectors and graph, chunk text and provenance, +document metadata, the next document id, construction settings, and embedder +identity. `/stats` reports whether persistence is enabled, whether state was loaded +at startup, and the service-state format version, without revealing the path. + +The Rust index uses a dependency-free little-endian binary format with an +`HNSWSNP` magic header, explicit version, dimensions/metric/construction parameters, +current SplitMix64 state, vectors, per-level graph links, and an FNV-1a checksum. +Loading rejects unsupported versions, truncation, checksum failures, non-finite +vectors, unreasonable sizes, and invalid graph references. Search-only scratch +buffers are reconstructed rather than serialized. The service wraps that binary +snapshot with deterministic JSON metadata in a versioned `RAGSTATE` container and +a BLAKE2b checksum. + +Each persistent insertion is staged against a cloned index while holding the store +lock. The complete container is written to a sibling temporary file, flushed and +`fsync`ed, atomically replaced, and its parent directory is `fsync`ed before the +new in-memory state is published and HTTP success is returned. A failed commit +therefore retains the previous file and previous live corpus. Corrupt state causes +startup to fail; it is never silently discarded. Stored index settings, embedding +dimension, backend class, and model name (where applicable) must match the runtime. + +Limitations: snapshots coordinate one service process only; do not point multiple +workers or hosts at the same file. Durability ultimately depends on the filesystem's +`fsync` and atomic-replace semantics. There is no deletion, migration between +embedder configurations, or automatic recovery of a corrupt snapshot. + ### Public API hardening | Variable | Default | Behavior | |----------|---------|----------| +| `RAG_STATE_PATH` | empty | Enables the versioned durable corpus snapshot at the configured file. Empty preserves in-memory behavior. | | `RAG_UPLOADS_ENABLED` | `0` | Enables `POST /documents` only when explicitly set to `1`, `true`, `yes`, or `on`. | | `RAG_MAX_TITLE_CHARS` | `200` | Maximum document title length. | | `RAG_MAX_DOCUMENT_CHARS` | `1000000` | Maximum document text length. | diff --git a/app/app/lib/rag.ts b/app/app/lib/rag.ts index aebde6d..a4f48bb 100644 --- a/app/app/lib/rag.ts +++ b/app/app/lib/rag.ts @@ -40,6 +40,7 @@ export interface Stats { embedder: string; generation: "mock" | "claude"; uploads_enabled: boolean; + persistence: { enabled: boolean; loaded: boolean; format_version: number }; limits: { title_chars: number; document_chars: number; diff --git a/bindings/python/hnsw_rag/embeddings.py b/bindings/python/hnsw_rag/embeddings.py index db797c5..fc142a9 100644 --- a/bindings/python/hnsw_rag/embeddings.py +++ b/bindings/python/hnsw_rag/embeddings.py @@ -71,6 +71,7 @@ class SentenceTransformerEmbedder: def __init__(self, model_name: str = "all-MiniLM-L6-v2") -> None: from sentence_transformers import SentenceTransformer # type: ignore + self.model_name = model_name self._model = SentenceTransformer(model_name) self.dim = int(self._model.get_sentence_embedding_dimension()) diff --git a/bindings/src/lib.rs b/bindings/src/lib.rs index faa8281..9125017 100644 --- a/bindings/src/lib.rs +++ b/bindings/src/lib.rs @@ -4,8 +4,9 @@ //! boundary as plain Python sequences of floats — no NumPy dependency, which //! keeps the module tiny; convert with `list(array)` if you have arrays. -use pyo3::exceptions::{PyIndexError, PyValueError}; +use pyo3::exceptions::{PyIOError, PyIndexError, PyValueError}; use pyo3::prelude::*; +use pyo3::types::PyBytes; use hnsw_engine as engine; @@ -26,6 +27,13 @@ fn parse_metric(metric: &str) -> PyResult { } } +fn snapshot_error(error: engine::Error) -> PyErr { + match error { + engine::Error::Io(_) => PyIOError::new_err(error.to_string()), + _ => PyValueError::new_err(error.to_string()), + } +} + /// An HNSW approximate-nearest-neighbor index over `dim`-dimensional vectors. #[pyclass] struct Hnsw { @@ -120,6 +128,30 @@ impl Hnsw { .ok_or_else(|| PyIndexError::new_err(format!("no vector with id {id}"))) } + /// Atomically save the versioned binary index snapshot. + fn save(&self, path: &str) -> PyResult<()> { + self.inner.save(path).map_err(snapshot_error) + } + + /// Restore an index, rejecting corrupt or unsupported snapshots. + #[staticmethod] + fn load(path: &str) -> PyResult { + engine::Hnsw::load(path) + .map(|inner| Self { inner }) + .map_err(snapshot_error) + } + + fn to_bytes<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> { + PyBytes::new(py, &self.inner.to_bytes()) + } + + #[staticmethod] + fn from_bytes(data: &[u8]) -> PyResult { + engine::Hnsw::from_bytes(data) + .map(|inner| Self { inner }) + .map_err(snapshot_error) + } + fn __len__(&self) -> usize { self.inner.len() } diff --git a/bindings/tests/test_hnsw.py b/bindings/tests/test_hnsw.py index 355609e..7f00a52 100644 --- a/bindings/tests/test_hnsw.py +++ b/bindings/tests/test_hnsw.py @@ -131,3 +131,20 @@ def test_recall_against_brute_force(): total += len(truth & got) / k recall = total / len(queries) assert recall >= 0.95, f"recall@{k} too low: {recall:.4f}" + + +def test_snapshot_roundtrip_and_continue_insert(tmp_path): + path = tmp_path / "index.hnsw" + index = Hnsw(dim=2, metric="euclidean", seed=9) + index.insert_batch([[0.0, 0.0], [2.0, 0.0]]) + expected = index.search([1.9, 0.0], k=2) + index.save(str(path)) + + restored = Hnsw.load(str(path)) + assert restored.search([1.9, 0.0], k=2) == pytest.approx(expected) + assert restored.insert([4.0, 0.0]) == 2 + assert restored.search([4.0, 0.0], k=1)[0][0] == 2 + + path.write_bytes(path.read_bytes()[:-3]) + with pytest.raises(ValueError, match="snapshot"): + Hnsw.load(str(path)) diff --git a/engine/src/hnsw.rs b/engine/src/hnsw.rs index 883b8ea..eaa24ed 100644 --- a/engine/src/hnsw.rs +++ b/engine/src/hnsw.rs @@ -16,6 +16,8 @@ use std::cmp::Reverse; use std::collections::BinaryHeap; +use std::io::{self, Write}; +use std::path::Path; use crate::distance::{self, Metric}; use crate::rng::SplitMix64; @@ -51,10 +53,32 @@ pub struct Neighbor { pub distance: f32, } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Debug)] pub enum Error { DimensionMismatch { expected: usize, got: usize }, NonFiniteValue { position: usize }, + Io(io::Error), + InvalidSnapshot(String), +} + +impl PartialEq for Error { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + ( + Self::DimensionMismatch { + expected: a, + got: b, + }, + Self::DimensionMismatch { + expected: c, + got: d, + }, + ) => a == c && b == d, + (Self::NonFiniteValue { position: a }, Self::NonFiniteValue { position: b }) => a == b, + (Self::InvalidSnapshot(a), Self::InvalidSnapshot(b)) => a == b, + _ => false, + } + } } impl std::fmt::Display for Error { @@ -72,12 +96,26 @@ impl std::fmt::Display for Error { "vector contains a non-finite value at position {position}" ) } + Error::Io(error) => write!(f, "snapshot I/O error: {error}"), + Error::InvalidSnapshot(message) => write!(f, "invalid HNSW snapshot: {message}"), } } } impl std::error::Error for Error {} +impl From for Error { + fn from(value: io::Error) -> Self { + Self::Io(value) + } +} + +pub const SNAPSHOT_VERSION: u32 = 1; +const SNAPSHOT_MAGIC: &[u8; 8] = b"HNSWSNP\0"; +const MAX_DIM: usize = 1_000_000; +const MAX_NODES: usize = u32::MAX as usize; +const MAX_LEVELS: usize = 128; + /// f32 wrapper with total order so distances can live in heaps. #[derive(Clone, Copy, PartialEq)] struct OrdF32(f32); @@ -220,6 +258,193 @@ impl Hnsw { self.store.metric } + pub fn params(&self) -> HnswParams { + self.params + } + + /// Encode all durable index state. Scratch search buffers are rebuilt. + pub fn to_bytes(&self) -> Vec { + let mut out = Vec::new(); + out.extend_from_slice(SNAPSHOT_MAGIC); + put_u32(&mut out, SNAPSHOT_VERSION); + out.push(match self.store.metric { + Metric::Euclidean => 0, + Metric::Cosine => 1, + }); + out.extend_from_slice(&[0; 3]); + put_u64(&mut out, self.store.dim as u64); + put_u64(&mut out, self.params.m as u64); + put_u64(&mut out, self.params.ef_construction as u64); + put_u64(&mut out, self.params.seed); + put_u64(&mut out, self.rng.state()); + put_u64(&mut out, self.store.len() as u64); + put_u64(&mut out, self.entry.map_or(u64::MAX, u64::from)); + put_u64(&mut out, self.top_level as u64); + for value in &self.store.data { + out.extend_from_slice(&value.to_bits().to_le_bytes()); + } + for node in &self.links { + put_u32(&mut out, node.len() as u32); + for level in node { + put_u32(&mut out, level.len() as u32); + for &id in level { + put_u32(&mut out, id); + } + } + } + let checksum = checksum(&out); + put_u64(&mut out, checksum); + out + } + + /// Decode a checked, versioned snapshot without unsafe deserialization. + pub fn from_bytes(bytes: &[u8]) -> Result { + if bytes.len() < 80 { + return invalid("snapshot is truncated"); + } + let (payload, trailer) = bytes.split_at(bytes.len() - 8); + let expected = u64::from_le_bytes(trailer.try_into().expect("eight-byte trailer")); + if checksum(payload) != expected { + return invalid("checksum mismatch"); + } + let mut r = Reader::new(payload); + if r.take(8)? != SNAPSHOT_MAGIC { + return invalid("bad magic header"); + } + let version = r.u32()?; + if version != SNAPSHOT_VERSION { + return invalid(format!("unsupported format version {version}")); + } + let metric = match r.u8()? { + 0 => Metric::Euclidean, + 1 => Metric::Cosine, + value => return invalid(format!("unknown metric tag {value}")), + }; + if r.take(3)? != [0, 0, 0] { + return invalid("non-zero reserved header bytes"); + } + let dim = usize_field(r.u64()?, "dimension", MAX_DIM)?; + let m = usize_field(r.u64()?, "m", MAX_NODES / 2)?; + let ef_construction = usize_field(r.u64()?, "ef_construction", MAX_NODES)?; + let seed = r.u64()?; + let rng_state = r.u64()?; + let count = usize_field(r.u64()?, "node count", MAX_NODES)?; + let entry_raw = r.u64()?; + let top_level = usize_field(r.u64()?, "top level", MAX_LEVELS - 1)?; + if dim == 0 || m < 2 || ef_construction == 0 { + return invalid("impossible index parameters"); + } + let values = count + .checked_mul(dim) + .ok_or_else(|| Error::InvalidSnapshot("vector count overflow".into()))?; + if values > r.remaining() / 4 { + return invalid("truncated vector data"); + } + let mut data = Vec::with_capacity(values); + for position in 0..values { + let value = f32::from_bits(r.u32()?); + if !value.is_finite() { + return invalid(format!("non-finite vector value at position {position}")); + } + data.push(value); + } + let mut links = Vec::with_capacity(count); + for node_id in 0..count { + let levels = usize_field(u64::from(r.u32()?), "level count", MAX_LEVELS)?; + if levels == 0 { + return invalid(format!("node {node_id} has no levels")); + } + let mut node = Vec::with_capacity(levels); + for level in 0..levels { + let cap = if level == 0 { + m.checked_mul(2) + .ok_or_else(|| Error::InvalidSnapshot("m overflow".into()))? + } else { + m + }; + let n = usize_field(u64::from(r.u32()?), "link count", cap)?; + let mut neighbors = Vec::with_capacity(n); + for _ in 0..n { + let id = r.u32()?; + if id as usize >= count { + return invalid(format!("node {node_id} references invalid node {id}")); + } + neighbors.push(id); + } + node.push(neighbors); + } + links.push(node); + } + if r.remaining() != 0 { + return invalid("trailing payload bytes"); + } + let entry = if entry_raw == u64::MAX { + None + } else { + Some( + u32::try_from(entry_raw) + .map_err(|_| Error::InvalidSnapshot("entry point is out of range".into()))?, + ) + }; + if count == 0 { + if entry.is_some() || top_level != 0 { + return invalid("empty index has an entry point or top level"); + } + } else { + let ep = entry.ok_or_else(|| { + Error::InvalidSnapshot("non-empty index has no entry point".into()) + })?; + if ep as usize >= count || links[ep as usize].len() != top_level + 1 { + return invalid("entry point/top level is inconsistent"); + } + if links.iter().any(|node| node.len() > top_level + 1) { + return invalid("node exceeds declared top level"); + } + } + Ok(Self { + params: HnswParams { + m, + ef_construction, + seed, + }, + store: VectorStore { dim, metric, data }, + links, + entry, + top_level, + m_max0: m * 2, + level_mult: 1.0 / (m as f64).ln(), + rng: SplitMix64::new(rng_state), + visited: Visited::new(), + }) + } + + pub fn save(&self, path: impl AsRef) -> Result<(), Error> { + let path = path.as_ref(); + let name = path + .file_name() + .and_then(|n| n.to_str()) + .ok_or_else(|| Error::InvalidSnapshot("snapshot path has no valid file name".into()))?; + let tmp = path.with_file_name(format!(".{name}.{}.tmp", std::process::id())); + let result = (|| -> Result<(), Error> { + let mut file = std::fs::File::create(&tmp)?; + file.write_all(&self.to_bytes())?; + file.sync_all()?; + std::fs::rename(&tmp, path)?; + if let Some(parent) = path.parent() { + std::fs::File::open(parent)?.sync_all()?; + } + Ok(()) + })(); + if result.is_err() { + let _ = std::fs::remove_file(tmp); + } + result + } + + pub fn load(path: impl AsRef) -> Result { + Self::from_bytes(&std::fs::read(path)?) + } + /// The stored vector for `id` (normalized if the metric is Cosine). pub fn vector(&self, id: u32) -> Option<&[f32]> { if (id as usize) < self.store.len() { @@ -364,6 +589,66 @@ impl Hnsw { } } +fn put_u32(out: &mut Vec, value: u32) { + out.extend_from_slice(&value.to_le_bytes()); +} +fn put_u64(out: &mut Vec, value: u64) { + out.extend_from_slice(&value.to_le_bytes()); +} +fn checksum(bytes: &[u8]) -> u64 { + bytes.iter().fold(0xcbf29ce484222325, |hash, byte| { + (hash ^ u64::from(*byte)).wrapping_mul(0x100000001b3) + }) +} +fn invalid(message: impl Into) -> Result { + Err(Error::InvalidSnapshot(message.into())) +} +fn usize_field(value: u64, name: &str, max: usize) -> Result { + let value = usize::try_from(value) + .map_err(|_| Error::InvalidSnapshot(format!("{name} does not fit this platform")))?; + if value > max { + return invalid(format!("{name} exceeds supported limit")); + } + Ok(value) +} +struct Reader<'a> { + bytes: &'a [u8], + pos: usize, +} +impl<'a> Reader<'a> { + fn new(bytes: &'a [u8]) -> Self { + Self { bytes, pos: 0 } + } + fn remaining(&self) -> usize { + self.bytes.len() - self.pos + } + fn take(&mut self, n: usize) -> Result<&'a [u8], Error> { + let end = self + .pos + .checked_add(n) + .ok_or_else(|| Error::InvalidSnapshot("offset overflow".into()))?; + if end > self.bytes.len() { + return invalid("snapshot is truncated"); + } + let value = &self.bytes[self.pos..end]; + self.pos = end; + Ok(value) + } + fn u8(&mut self) -> Result { + Ok(self.take(1)?[0]) + } + fn u32(&mut self) -> Result { + Ok(u32::from_le_bytes( + self.take(4)?.try_into().expect("four bytes"), + )) + } + fn u64(&mut self) -> Result { + Ok(u64::from_le_bytes( + self.take(8)?.try_into().expect("eight bytes"), + )) + } +} + /// Algorithm 2 from the paper: beam search within one layer. Takes entry /// points as `(comparison_distance, id)` pairs, returns up to `ef` closest /// nodes, sorted ascending by distance. diff --git a/engine/src/lib.rs b/engine/src/lib.rs index 32e7244..48d71cb 100644 --- a/engine/src/lib.rs +++ b/engine/src/lib.rs @@ -21,4 +21,4 @@ mod hnsw; pub mod rng; pub use distance::Metric; -pub use hnsw::{Error, Hnsw, HnswParams, Neighbor}; +pub use hnsw::{Error, Hnsw, HnswParams, Neighbor, SNAPSHOT_VERSION}; diff --git a/engine/src/rng.rs b/engine/src/rng.rs index 21aaf24..0c9b0ea 100644 --- a/engine/src/rng.rs +++ b/engine/src/rng.rs @@ -14,6 +14,10 @@ impl SplitMix64 { Self { state: seed } } + pub(crate) fn state(&self) -> u64 { + self.state + } + pub fn next_u64(&mut self) -> u64 { self.state = self.state.wrapping_add(0x9E37_79B9_7F4A_7C15); let mut z = self.state; diff --git a/engine/tests/snapshot.rs b/engine/tests/snapshot.rs new file mode 100644 index 0000000..e05b377 --- /dev/null +++ b/engine/tests/snapshot.rs @@ -0,0 +1,83 @@ +use hnsw_engine::{Hnsw, HnswParams, Metric, SNAPSHOT_VERSION}; + +fn params() -> HnswParams { + HnswParams { + m: 4, + ef_construction: 32, + seed: 17, + } +} +fn checksum(bytes: &[u8]) -> u64 { + bytes.iter().fold(0xcbf29ce484222325, |h, b| { + (h ^ u64::from(*b)).wrapping_mul(0x100000001b3) + }) +} + +#[test] +fn empty_snapshot_round_trip() { + let index = Hnsw::new(3, Metric::Cosine, params()); + let restored = Hnsw::from_bytes(&index.to_bytes()).unwrap(); + assert_eq!(restored.len(), 0); + assert_eq!(restored.dim(), 3); + assert_eq!(restored.metric(), Metric::Cosine); + assert_eq!(restored.params().m, 4); +} + +#[test] +fn populated_metrics_round_trip_with_identical_searches() { + for metric in [Metric::Cosine, Metric::Euclidean] { + let mut index = Hnsw::new(3, metric, params()); + for vector in [[1., 0., 0.], [0., 1., 0.], [0., 0., 1.], [1., 1., 0.]] { + index.insert(vector.to_vec()).unwrap(); + } + let before = index.search(&[0.9, 0.2, 0.], 4, 20).unwrap(); + let mut restored = Hnsw::from_bytes(&index.to_bytes()).unwrap(); + assert_eq!(restored.len(), index.len()); + assert_eq!(restored.dim(), index.dim()); + assert_eq!(restored.metric(), metric); + assert_eq!(restored.params().ef_construction, params().ef_construction); + assert_eq!(restored.search(&[0.9, 0.2, 0.], 4, 20).unwrap(), before); + assert_eq!(restored.insert(vec![-1., 0., 0.]).unwrap(), 4); + assert_eq!(restored.search(&[-1., 0., 0.], 1, 20).unwrap()[0].id, 4); + } +} + +#[test] +fn file_api_and_rng_continuation_are_exact() { + let mut original = Hnsw::new(2, Metric::Euclidean, params()); + original.insert(vec![0., 0.]).unwrap(); + let path = std::env::temp_dir().join(format!("hnsw-snapshot-{}.bin", std::process::id())); + original.save(&path).unwrap(); + let mut restored = Hnsw::load(&path).unwrap(); + std::fs::remove_file(path).unwrap(); + for vector in [[1., 0.], [2., 0.], [3., 0.]] { + original.insert(vector.to_vec()).unwrap(); + restored.insert(vector.to_vec()).unwrap(); + } + assert_eq!(original.to_bytes(), restored.to_bytes()); +} + +#[test] +fn corrupt_truncated_and_unsupported_snapshots_fail() { + let bytes = Hnsw::new(2, Metric::Cosine, params()).to_bytes(); + for end in 0..bytes.len() { + assert!(Hnsw::from_bytes(&bytes[..end]).is_err()); + } + let mut corrupt = bytes.clone(); + corrupt[20] ^= 1; + assert!(Hnsw::from_bytes(&corrupt) + .err() + .unwrap() + .to_string() + .contains("checksum")); + let mut unsupported = bytes; + unsupported[8..12].copy_from_slice(&(SNAPSHOT_VERSION + 1).to_le_bytes()); + let len = unsupported.len(); + let sum = checksum(&unsupported[..len - 8]); + unsupported[len - 8..].copy_from_slice(&sum.to_le_bytes()); + assert!(Hnsw::from_bytes(&unsupported) + .err() + .unwrap() + .to_string() + .contains("unsupported")); +} diff --git a/service/rag_service/app.py b/service/rag_service/app.py index 766f04e..12b4464 100644 --- a/service/rag_service/app.py +++ b/service/rag_service/app.py @@ -154,12 +154,13 @@ async def _reject(self, scope: Scope, receive: Receive, send: Send) -> None: await response(scope, receive, send) -# One in-memory store for the process. The embedder backend is chosen at +# One process-local store; durability is opt-in through RAG_STATE_PATH. The embedder backend is chosen at # startup: real model if available, deterministic hashed fallback otherwise. _embedder = get_embedder(os.environ.get("RAG_EMBEDDER", "auto")) _store = DocumentStore( embedder=_embedder, min_score=_finite_env_float("RAG_MIN_SCORE", 0.09), + state_path=os.environ.get("RAG_STATE_PATH"), ) SAMPLE_DOCS = pathlib.Path(__file__).resolve().parent.parent / "sample_docs" @@ -168,9 +169,7 @@ async def _reject(self, scope: Scope, receive: Receive, send: Send) -> None: def seed_sample_docs() -> int: """Index the bundled sample corpus in-process. Returns documents added. - The index is in-memory, so a fresh container starts empty. `scripts/seed.py` - solves that over HTTP for a running instance, but a hosted deployment has - nobody to run it — hence this in-process path for startup. + The hosted demo has persistence disabled, so startup seeding remains useful. """ if not SAMPLE_DOCS.is_dir(): log.warning("sample corpus not found at %s; starting with an empty index", SAMPLE_DOCS) diff --git a/service/rag_service/store.py b/service/rag_service/store.py index cbd2cf4..ff560f0 100644 --- a/service/rag_service/store.py +++ b/service/rag_service/store.py @@ -1,25 +1,28 @@ -"""In-memory document + chunk store backed by the HNSW index. - -Keeps three things in lockstep: -- the HNSW index (chunk vectors, keyed by dense id = insertion order), -- `chunks[id]` -> the chunk's text + provenance, -- `documents` -> per-document metadata. - -The index is created lazily on the first insert, because its dimension is -whatever the configured embedder produces. Everything lives in memory; this -is a demo/portfolio service, not a durable database, and it says so. -""" +"""Thread-safe document store with optional durable, atomic snapshots.""" from __future__ import annotations +import hashlib +import json +import os +import pathlib +import struct import threading -from dataclasses import dataclass, field +from dataclasses import asdict, dataclass, field from typing import Dict, List, Optional from hnsw_engine import Hnsw from hnsw_rag import Chunk, chunk_text from hnsw_rag.embeddings import Embedder +STATE_MAGIC = b"RAGSTATE" +STATE_VERSION = 1 +HNSW_FORMAT_VERSION = 1 +_HEADER = struct.Struct("<8sIQQ") +_CHECKSUM_BYTES = 32 +_MAX_METADATA_BYTES = 256 * 1024 * 1024 +_MAX_INDEX_BYTES = 16 * 1024 * 1024 * 1024 + @dataclass class StoredChunk: @@ -27,7 +30,7 @@ class StoredChunk: text: str doc_id: int doc_title: str - ordinal: int # position within the source document + ordinal: int @dataclass @@ -44,7 +47,15 @@ class RetrievedChunk: doc_id: int doc_title: str ordinal: int - score: float # cosine similarity in [~-1, 1]; higher = more relevant + score: float + + +def _embedder_identity(embedder: Embedder) -> dict: + identity = {"type": type(embedder).__name__, "dim": embedder.dim} + model_name = getattr(embedder, "model_name", None) + if model_name is not None: + identity["model_name"] = model_name + return identity @dataclass @@ -55,97 +66,242 @@ class DocumentStore: m: int = 16 ef_construction: int = 200 seed: int = 0x5EED + state_path: Optional[pathlib.Path | str] = None _index: Optional[Hnsw] = None _chunks: Dict[int, StoredChunk] = field(default_factory=dict) _documents: Dict[int, Document] = field(default_factory=dict) _next_doc_id: int = 0 + _loaded: bool = False _lock: threading.Lock = field(default_factory=threading.Lock) + def __post_init__(self) -> None: + if self.state_path is not None: + raw = str(self.state_path).strip() + self.state_path = pathlib.Path(raw) if raw else None + if self.state_path is not None and self.state_path.exists(): + self._load_state(self.state_path) + self._loaded = True + @property def dim(self) -> int: return self.embedder.dim + def _new_index(self) -> Hnsw: + return Hnsw( + dim=self.dim, + metric=self.metric, + m=self.m, + ef_construction=self.ef_construction, + seed=self.seed, + ) + + def _metadata( + self, + chunks: Dict[int, StoredChunk], + documents: Dict[int, Document], + next_doc_id: int, + ) -> dict: + return { + "state_version": STATE_VERSION, + "hnsw_format_version": HNSW_FORMAT_VERSION, + "index": { + "dim": self.dim, + "metric": self.metric, + "m": self.m, + "ef_construction": self.ef_construction, + "seed": self.seed, + }, + "embedder": _embedder_identity(self.embedder), + "next_doc_id": next_doc_id, + "chunks": [asdict(chunks[key]) for key in sorted(chunks)], + "documents": [asdict(documents[key]) for key in sorted(documents)], + } + + def _persist( + self, + index: Optional[Hnsw], + chunks: Dict[int, StoredChunk], + documents: Dict[int, Document], + next_doc_id: int, + ) -> None: + assert self.state_path is not None + actual_index = index if index is not None else self._new_index() + metadata = json.dumps( + self._metadata(chunks, documents, next_doc_id), + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8") + index_data = bytes(actual_index.to_bytes()) + payload = ( + _HEADER.pack(STATE_MAGIC, STATE_VERSION, len(metadata), len(index_data)) + + metadata + + index_data + ) + data = payload + hashlib.blake2b(payload, digest_size=_CHECKSUM_BYTES).digest() + path = self.state_path + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name( + f".{path.name}.{os.getpid()}.{threading.get_ident()}.tmp" + ) + try: + with temporary.open("xb") as output: + output.write(data) + output.flush() + os.fsync(output.fileno()) + os.replace(temporary, path) + directory_fd = os.open(path.parent, os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + finally: + try: + temporary.unlink() + except FileNotFoundError: + pass + + def _load_state(self, path: pathlib.Path) -> None: + data = path.read_bytes() + if len(data) < _HEADER.size + _CHECKSUM_BYTES: + raise RuntimeError("persisted RAG state is truncated") + payload, digest = data[:-_CHECKSUM_BYTES], data[-_CHECKSUM_BYTES:] + if hashlib.blake2b(payload, digest_size=_CHECKSUM_BYTES).digest() != digest: + raise RuntimeError("persisted RAG state checksum mismatch") + magic, version, metadata_len, index_len = _HEADER.unpack_from(payload) + if magic != STATE_MAGIC: + raise RuntimeError("persisted RAG state has an invalid magic header") + if version != STATE_VERSION: + raise RuntimeError(f"unsupported RAG state format version {version}") + if metadata_len > _MAX_METADATA_BYTES or index_len > _MAX_INDEX_BYTES: + raise RuntimeError("persisted RAG state exceeds supported size limits") + expected = _HEADER.size + metadata_len + index_len + if expected != len(payload): + raise RuntimeError("persisted RAG state lengths are invalid") + metadata_end = _HEADER.size + metadata_len + try: + metadata = json.loads(payload[_HEADER.size : metadata_end]) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise RuntimeError("persisted RAG metadata is invalid") from exc + expected_index = { + "dim": self.dim, + "metric": self.metric, + "m": self.m, + "ef_construction": self.ef_construction, + "seed": self.seed, + } + if metadata.get("index") != expected_index: + raise RuntimeError( + f"persisted index configuration is incompatible with runtime configuration (stored={metadata.get('index')!r}, runtime={expected_index!r})" + ) + identity = _embedder_identity(self.embedder) + if metadata.get("embedder") != identity: + raise RuntimeError( + f"persisted embedder is incompatible with runtime embedder (stored={metadata.get('embedder')!r}, runtime={identity!r})" + ) + try: + index = Hnsw.from_bytes(payload[metadata_end:]) + chunks = {item["id"]: StoredChunk(**item) for item in metadata["chunks"]} + documents = {item["id"]: Document(**item) for item in metadata["documents"]} + next_doc_id = int(metadata["next_doc_id"]) + except (KeyError, TypeError, ValueError) as exc: + raise RuntimeError("persisted RAG metadata has an invalid schema") from exc + if ( + metadata.get("state_version") != STATE_VERSION + or metadata.get("hnsw_format_version") != HNSW_FORMAT_VERSION + ): + raise RuntimeError("persisted RAG metadata version is incompatible") + if index.dim != self.dim or index.metric != self.metric: + raise RuntimeError( + "persisted HNSW configuration does not match its metadata" + ) + if len(index) != len(chunks) or set(chunks) != set(range(len(index))): + raise RuntimeError("persisted HNSW ids and chunk metadata are not aligned") + if next_doc_id < 0 or any(doc_id >= next_doc_id for doc_id in documents): + raise RuntimeError("persisted document insertion state is invalid") + if any(chunk.doc_id not in documents for chunk in chunks.values()): + raise RuntimeError("persisted chunk references an unknown document") + chunk_counts = {doc_id: 0 for doc_id in documents} + for chunk in chunks.values(): + chunk_counts[chunk.doc_id] += 1 + if chunk.doc_title != documents[chunk.doc_id].title: + raise RuntimeError("persisted chunk and document titles disagree") + if any( + document.n_chunks != chunk_counts[doc_id] + for doc_id, document in documents.items() + ): + raise RuntimeError("persisted document chunk counts are invalid") + self._index, self._chunks, self._documents, self._next_doc_id = ( + index, + chunks, + documents, + next_doc_id, + ) + def add_document( self, title: str, text: str, *, max_words: int = 180, overlap: int = 40 ) -> Document: - """Chunk, embed, and index a document. Returns its metadata.""" chunks: List[Chunk] = chunk_text( text, max_words=max_words, overlap=overlap, source=title ) - if not chunks: - with self._lock: - doc_id = self._next_doc_id - self._next_doc_id += 1 - doc = Document(id=doc_id, title=title, n_chunks=0) - self._documents[doc_id] = doc - return doc - - vectors = self.embedder.embed([c.text for c in chunks]) + vectors = self.embedder.embed([c.text for c in chunks]) if chunks else [] if len(vectors) != len(chunks): raise ValueError( f"embedder returned {len(vectors)} vectors for {len(chunks)} chunks" ) - with self._lock: - if self._index is None: - self._index = Hnsw( - dim=self.dim, - metric=self.metric, - m=self.m, - ef_construction=self.ef_construction, - seed=self.seed, + # With persistence enabled, mutate a snapshot clone, commit it, then + # publish it in memory. A failed fsync/replace leaves both old states active. + staging = self.state_path is not None + index = ( + Hnsw.from_bytes(bytes(self._index.to_bytes())) + if staging and self._index is not None + else ( + self._new_index() if chunks and self._index is None else self._index ) + ) + chunk_map = dict(self._chunks) if staging else self._chunks + document_map = dict(self._documents) if staging else self._documents doc_id = self._next_doc_id - self._next_doc_id += 1 - ids = self._index.insert_batch(vectors) + next_doc_id = doc_id + 1 + ids = index.insert_batch(vectors) if chunks and index is not None else [] if len(ids) != len(chunks): raise RuntimeError( f"index returned {len(ids)} ids for {len(chunks)} chunks" ) for chunk, cid in zip(chunks, ids): - self._chunks[cid] = StoredChunk( - id=cid, - text=chunk.text, - doc_id=doc_id, - doc_title=title, - ordinal=chunk.index, + chunk_map[cid] = StoredChunk( + cid, chunk.text, doc_id, title, chunk.index ) - doc = Document(id=doc_id, title=title, n_chunks=len(chunks)) - self._documents[doc_id] = doc + doc = Document(doc_id, title, len(chunks)) + document_map[doc_id] = doc + if staging: + self._persist(index, chunk_map, document_map, next_doc_id) + self._index, self._chunks, self._documents = index, chunk_map, document_map + self._next_doc_id = next_doc_id return doc - def retrieve(self, query: str, k: int = 5, ef_search: int = 100) -> List[RetrievedChunk]: - """Embed the query and return the k most relevant chunks.""" + def retrieve( + self, query: str, k: int = 5, ef_search: int = 100 + ) -> List[RetrievedChunk]: qvec = self.embedder.embed([query])[0] with self._lock: if self._index is None or len(self._index) == 0: return [] - # The Python wrapper exposes insert as a mutable operation. Keep - # search and its metadata lookup in the same critical section so - # concurrent uploads cannot mutate the graph beneath a query. hits = self._index.search(qvec, k=k, ef_search=ef_search) - out: List[RetrievedChunk] = [] + out = [] for cid, distance in hits: sc = self._chunks[cid] - # Cosine distance is 1 - similarity; report similarity so higher - # is more relevant, which is what a reader expects from a score. score = 1.0 - distance if self.metric == "cosine" else -distance if self.min_score is not None and score < self.min_score: - # Hits are closest-first, so their relevance scores only - # decrease. Weak nearest neighbors are not useful grounding. break out.append( RetrievedChunk( - id=sc.id, - text=sc.text, - doc_id=sc.doc_id, - doc_title=sc.doc_title, - ordinal=sc.ordinal, - score=score, + sc.id, sc.text, sc.doc_id, sc.doc_title, sc.ordinal, score ) ) - return out + return out def stats(self) -> dict: with self._lock: @@ -157,10 +313,12 @@ def stats(self) -> dict: "min_score": self.min_score, "m": self.m, "ef_construction": self.ef_construction, - # Which embedding backend is actually live. Worth surfacing: - # HashedEmbedder matches on term overlap, not meaning, so a - # reader should not mistake it for semantic search. "embedder": type(self.embedder).__name__, + "persistence": { + "enabled": self.state_path is not None, + "loaded": self._loaded, + "format_version": STATE_VERSION, + }, } def documents(self) -> List[Document]: diff --git a/service/tests/test_e2e.py b/service/tests/test_e2e.py index eca9045..b5b71b3 100644 --- a/service/tests/test_e2e.py +++ b/service/tests/test_e2e.py @@ -61,6 +61,11 @@ def test_stats_after_seed(client): assert stats["embedder"] == "HashedEmbedder" assert stats["generation"] == "mock" assert stats["uploads_enabled"] is False + assert stats["persistence"] == { + "enabled": False, + "loaded": False, + "format_version": 1, + } assert stats["limits"] == { "title_chars": MAX_TITLE_CHARS, "document_chars": MAX_DOCUMENT_CHARS, diff --git a/service/tests/test_persistence.py b/service/tests/test_persistence.py new file mode 100644 index 0000000..b9d6e4c --- /dev/null +++ b/service/tests/test_persistence.py @@ -0,0 +1,51 @@ +import pytest + +from hnsw_rag.embeddings import HashedEmbedder +from rag_service.store import DocumentStore, STATE_VERSION + + +def test_complete_store_survives_restart(tmp_path): + path = tmp_path / "corpus.state" + first = DocumentStore(HashedEmbedder(dim=32), state_path=path, min_score=None) + doc = first.add_document("durable title", "persistent graph metadata provenance") + before = first.retrieve("persistent metadata", k=1) + + second = DocumentStore(HashedEmbedder(dim=32), state_path=path, min_score=None) + after = second.retrieve("persistent metadata", k=1) + assert second.documents() == [doc] + assert [(x.id, x.text, x.doc_id, x.doc_title, x.ordinal) for x in after] == [ + (x.id, x.text, x.doc_id, x.doc_title, x.ordinal) for x in before + ] + assert second.stats()["persistence"] == { + "enabled": True, + "loaded": True, + "format_version": STATE_VERSION, + } + assert second.add_document("next", "another durable chunk").id == 1 + + +def test_corrupt_and_incompatible_state_rejected(tmp_path): + path = tmp_path / "corpus.state" + DocumentStore(HashedEmbedder(dim=16), state_path=path).add_document( + "one", "content" + ) + with pytest.raises(RuntimeError, match="incompatible"): + DocumentStore(HashedEmbedder(dim=8), state_path=path) + data = bytearray(path.read_bytes()) + data[-1] ^= 1 + path.write_bytes(data) + with pytest.raises(RuntimeError, match="checksum"): + DocumentStore(HashedEmbedder(dim=16), state_path=path) + + +def test_failed_commit_does_not_publish_staged_mutation(tmp_path, monkeypatch): + store = DocumentStore(HashedEmbedder(dim=8), state_path=tmp_path / "state") + store.add_document("first", "one") + before = store.stats() + monkeypatch.setattr( + store, "_persist", lambda *args: (_ for _ in ()).throw(OSError("disk full")) + ) + with pytest.raises(OSError, match="disk full"): + store.add_document("second", "two") + assert store.stats()["documents"] == before["documents"] + assert store.stats()["chunks"] == before["chunks"]