diff --git a/Cargo.lock b/Cargo.lock index e8ccb45a..e98615a9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -321,34 +321,13 @@ dependencies = [ "crypto-common 0.2.2", ] -[[package]] -name = "dirs" -version = "5.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44c45a9d03d6676652bcb5e724c7e988de1acad23a711b5217ab9cbecbec2225" -dependencies = [ - "dirs-sys 0.4.1", -] - [[package]] name = "dirs" version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3e8aa94d75141228480295a7d0e7feb620b1a5ad9f12bc40be62411e38cce4e" dependencies = [ - "dirs-sys 0.5.0", -] - -[[package]] -name = "dirs-sys" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "520f05a5cbd335fae5a99ff7a6ab8627577660ee5cfd6a94a6a929b52ff0321c" -dependencies = [ - "libc", - "option-ext", - "redox_users 0.4.6", - "windows-sys 0.48.0", + "dirs-sys", ] [[package]] @@ -359,7 +338,7 @@ checksum = "e01a3366d27ee9890022452ee61b2b63a67e6f13f58900b651ff5665f0bb1fab" dependencies = [ "libc", "option-ext", - "redox_users 0.5.2", + "redox_users", "windows-sys 0.61.2", ] @@ -1168,7 +1147,7 @@ dependencies = [ "rustc-hash", "rustls", "socket2", - "thiserror 2.0.20", + "thiserror", "tokio", "tracing", "web-time", @@ -1190,7 +1169,7 @@ dependencies = [ "rustls", "rustls-pki-types", "slab", - "thiserror 2.0.20", + "thiserror", "tinyvec", "tracing", "web-time", @@ -1295,17 +1274,6 @@ dependencies = [ "bitflags", ] -[[package]] -name = "redox_users" -version = "0.4.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba009ff324d1fc1b900bd1fdb31564febe58a8ccc8a6fdbb93b543d33b13ca43" -dependencies = [ - "getrandom 0.2.17", - "libredox", - "thiserror 1.0.69", -] - [[package]] name = "redox_users" version = "0.5.2" @@ -1314,7 +1282,7 @@ checksum = "a4e608c6638b9c18977b00b475ac1f28d14e84b27d8d42f70e0bf1e3dec127ac" dependencies = [ "getrandom 0.2.17", "libredox", - "thiserror 2.0.20", + "thiserror", ] [[package]] @@ -1429,7 +1397,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c51c9ae4df8a7fba42103df5c621fa3c37eccf3a3c650879e90fc48b11cc192c" dependencies = [ "hashbrown 0.16.1", - "thiserror 2.0.20", + "thiserror", ] [[package]] @@ -1772,39 +1740,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.4.3", + "getrandom 0.3.4", "once_cell", "rustix", "windows-sys 0.61.2", ] -[[package]] -name = "thiserror" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" -dependencies = [ - "thiserror-impl 1.0.69", -] - [[package]] name = "thiserror" version = "2.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" dependencies = [ - "thiserror-impl 2.0.20", -] - -[[package]] -name = "thiserror-impl" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.119", + "thiserror-impl", ] [[package]] @@ -1820,13 +1768,13 @@ dependencies = [ [[package]] name = "tinycortex" -version = "0.1.1" +version = "0.1.2" dependencies = [ "anyhow", "async-trait", "block2", "chrono", - "dirs 5.0.1", + "dirs", "futures", "git2", "hex", @@ -1843,7 +1791,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.10.9", - "thiserror 2.0.20", + "thiserror", "tinycortex-api", "tinyinference", "tokio", @@ -1872,7 +1820,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.11.0", - "thiserror 2.0.20", + "thiserror", "tokio", ] @@ -1920,7 +1868,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.11.0", - "thiserror 2.0.20", + "thiserror", "uuid", ] @@ -1943,7 +1891,7 @@ dependencies = [ "async-trait", "axum", "chrono", - "dirs 6.0.0", + "dirs", "futures", "log", "parking_lot", @@ -1955,7 +1903,7 @@ dependencies = [ "serde_json", "sha2 0.11.0", "tempfile", - "thiserror 2.0.20", + "thiserror", "tinycortex", "tinycortex-api", "tinyinference", @@ -2317,7 +2265,7 @@ dependencies = [ "log", "rand 0.9.5", "sha1", - "thiserror 2.0.20", + "thiserror", ] [[package]] @@ -2606,22 +2554,13 @@ dependencies = [ "windows-link", ] -[[package]] -name = "windows-sys" -version = "0.48.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9" -dependencies = [ - "windows-targets 0.48.5", -] - [[package]] name = "windows-sys" version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets 0.52.6", + "windows-targets", ] [[package]] @@ -2633,67 +2572,34 @@ dependencies = [ "windows-link", ] -[[package]] -name = "windows-targets" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" -dependencies = [ - "windows_aarch64_gnullvm 0.48.5", - "windows_aarch64_msvc 0.48.5", - "windows_i686_gnu 0.48.5", - "windows_i686_msvc 0.48.5", - "windows_x86_64_gnu 0.48.5", - "windows_x86_64_gnullvm 0.48.5", - "windows_x86_64_msvc 0.48.5", -] - [[package]] name = "windows-targets" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ - "windows_aarch64_gnullvm 0.52.6", - "windows_aarch64_msvc 0.52.6", - "windows_i686_gnu 0.52.6", + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", "windows_i686_gnullvm", - "windows_i686_msvc 0.52.6", - "windows_x86_64_gnu 0.52.6", - "windows_x86_64_gnullvm 0.52.6", - "windows_x86_64_msvc 0.52.6", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", ] -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" - [[package]] name = "windows_aarch64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" -[[package]] -name = "windows_aarch64_msvc" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" - [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" -[[package]] -name = "windows_i686_gnu" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" - [[package]] name = "windows_i686_gnu" version = "0.52.6" @@ -2706,48 +2612,24 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" -[[package]] -name = "windows_i686_msvc" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" - [[package]] name = "windows_i686_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" -[[package]] -name = "windows_x86_64_gnu" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" - [[package]] name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" - [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" -[[package]] -name = "windows_x86_64_msvc" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" - [[package]] name = "windows_x86_64_msvc" version = "0.52.6" diff --git a/crates/tinymemory-core/src/store/client.rs b/crates/tinymemory-core/src/store/client.rs index e4808279..d82f720f 100644 --- a/crates/tinymemory-core/src/store/client.rs +++ b/crates/tinymemory-core/src/store/client.rs @@ -27,6 +27,23 @@ use tinymemory_api::host::EmbeddingProvider; /// Reference-counted handle to a `MemoryClient`. pub type MemoryClientRef = Arc; +/// Outcome of [`MemoryClient::put_docs`]. +#[derive(Debug, Default)] +pub struct BatchPutOutcome { + /// One entry per document attempted, in input order. The first failure + /// ends the batch and is always the last entry, so every document before + /// it was written and queued. + pub results: Vec>, + /// Written documents whose background graph-extraction job the ingestion + /// queue refused (it was full, or its worker had shut down). Those + /// documents are stored and searchable; only their entity/relation + /// extraction was skipped — the same best-effort drop a refused + /// [`MemoryClient::put_doc`] submission makes, counted here because a + /// batch submits hundreds of jobs in one burst where the per-document + /// write path spaced them out. + pub dropped_extractions: usize, +} + /// Thread-safe container for an optional `MemoryClientRef`. /// /// Used for global state management where the memory client may or may not @@ -186,6 +203,50 @@ impl MemoryClient { Ok(document_id) } + /// Store many documents at once — the batch form of [`Self::put_doc`]. + /// + /// The documents' chunks are embedded together, one provider request per + /// bounded group of chunk texts across the whole batch, instead of one + /// request per document (tinymemory#138). Each document is still gated, + /// written and queued for background graph extraction exactly as + /// `put_doc` does it. + /// + /// Documents are written in order and the first failure ends the batch: + /// [`BatchPutOutcome::results`] holds one entry per document attempted, in + /// input order, so a failure is always the last entry and every document + /// before it was written and queued. A graph-extraction job the ingestion + /// queue refuses is not a document failure — the document is stored — but + /// it is counted in [`BatchPutOutcome::dropped_extractions`] and logged, + /// so a burst that overruns the queue is visible to the caller. + pub async fn put_docs(&self, inputs: Vec) -> BatchPutOutcome { + let results = self.inner.upsert_documents(inputs.clone()).await; + let mut dropped_extractions = 0; + for (document, result) in inputs.into_iter().zip(&results) { + if let Ok(document_id) = result { + let queued = self.ingestion_queue.submit(IngestionJob { + document_id: document_id.clone(), + document, + config: MemoryIngestionConfig::default(), + }); + if !queued { + dropped_extractions += 1; + } + } + } + if dropped_extractions > 0 { + log::warn!( + "[memory] graph extraction skipped for {dropped_extractions} of {} written \ + document(s): the ingestion queue refused the job(s); the documents \ + themselves are stored", + results.iter().filter(|result| result.is_ok()).count() + ); + } + BatchPutOutcome { + results, + dropped_extractions, + } + } + /// Store a document (DB row + markdown file) without vector embedding or /// graph extraction. Use this for high-frequency, ephemeral writes where /// the full pipeline would be too expensive (e.g. transient sync diff --git a/crates/tinymemory-core/src/store/client_tests.rs b/crates/tinymemory-core/src/store/client_tests.rs index 20f6aabf..e41a994f 100644 --- a/crates/tinymemory-core/src/store/client_tests.rs +++ b/crates/tinymemory-core/src/store/client_tests.rs @@ -379,3 +379,97 @@ async fn ingest_doc_completes_and_stores_document() { // is exercised (no panic). let _ = result; } + +#[tokio::test] +async fn put_docs_writes_every_document_and_returns_ids_in_order() { + let (_tmp, client) = make_client(); + let outcome = client + .put_docs(vec![ + doc("batch", "k1", "one"), + doc("batch", "k2", "two"), + doc("batch", "k3", "three"), + ]) + .await; + assert_eq!( + outcome.dropped_extractions, 0, + "an idle default-capacity queue accepts every graph-extraction job" + ); + + let ids: Vec = outcome + .results + .into_iter() + .map(|result| result.expect("each document is written")) + .collect(); + assert_eq!(ids.len(), 3); + + let listed = client.list_documents(Some("batch")).await.unwrap(); + assert_eq!(listed["count"].as_u64(), Some(3)); + let by_key: std::collections::BTreeMap = listed["documents"] + .as_array() + .unwrap() + .iter() + .map(|document| { + ( + document["key"].as_str().unwrap().to_string(), + document["documentId"].as_str().unwrap().to_string(), + ) + }) + .collect(); + assert_eq!( + ids, + vec![ + by_key["k1"].clone(), + by_key["k2"].clone(), + by_key["k3"].clone() + ], + "ids come back in input order" + ); +} + +/// A batch submits its graph-extraction jobs in one burst, so a queue that a +/// per-document trickle never overran can refuse some of them. The refusal is +/// the queue's documented best-effort drop, but it must be counted rather +/// than lost: the caller sees how many documents skipped extraction. +#[tokio::test] +async fn put_docs_counts_the_graph_jobs_a_full_ingestion_queue_refuses() { + use tinymemory_api::host::NoopEmbedding; + crate::test_seams::init(); + let tmp = TempDir::new().unwrap(); + let inner = Arc::new(UnifiedMemory::new(tmp.path(), Arc::new(NoopEmbedding), None).unwrap()); + let state = IngestionState::new(); + // A one-slot queue whose worker cannot drain: the test holds the singleton + // run lock the worker takes before it processes a job, so after the worker + // pulls the first job the slot refills once and every later job is refused. + let ingestion_queue = + ingestion_queue::start_worker_with_capacity(Arc::clone(&inner), state.clone(), 1); + let _worker_blocked = state.acquire().await; + let client = MemoryClient { + inner, + ingestion_queue, + }; + + let outcome = client + .put_docs(vec![ + doc("burst", "k1", "one"), + doc("burst", "k2", "two"), + doc("burst", "k3", "three"), + ]) + .await; + + assert!( + outcome.results.iter().all(Result::is_ok), + "a refused extraction job is not a document failure, got {:?}", + outcome.results + ); + assert!( + (1..=2).contains(&outcome.dropped_extractions), + "one job fits the single slot (two if the worker pulled the first before the \ + next submit); the rest are refused and counted, got {}", + outcome.dropped_extractions + ); + assert_eq!( + client.list_documents(Some("burst")).await.unwrap()["count"].as_u64(), + Some(3), + "refused extraction jobs do not affect the documents themselves" + ); +} diff --git a/crates/tinymemory-core/src/store/mod.rs b/crates/tinymemory-core/src/store/mod.rs index d4aa5b9a..656b6279 100644 --- a/crates/tinymemory-core/src/store/mod.rs +++ b/crates/tinymemory-core/src/store/mod.rs @@ -55,7 +55,7 @@ mod write_gate; pub use kinds::MemoryKind; pub use traits::{ObsidianFile, ObsidianRepresentable, VectorEmbeddable}; -pub use client::{MemoryClient, MemoryClientRef, MemoryState}; +pub use client::{BatchPutOutcome, MemoryClient, MemoryClientRef, MemoryState}; pub use factories::{ active_embedding_signature, create_memory, create_memory_for_migration, create_memory_with_local_ai, effective_embedding_settings, effective_memory_backend_name, diff --git a/crates/tinymemory-core/src/store/namespace_store/README.md b/crates/tinymemory-core/src/store/namespace_store/README.md index 1718dcfb..d64f572f 100644 --- a/crates/tinymemory-core/src/store/namespace_store/README.md +++ b/crates/tinymemory-core/src/store/namespace_store/README.md @@ -11,7 +11,7 @@ tier; this directory is intentionally not migration staging. - **`mod.rs`** — declares the `UnifiedMemory` struct (connection + paths + embedder) and wires the submodules. - **`init.rs`** — constructor, `CREATE TABLE` bootstrap (docs, kv, graph, vector chunks, episodic FTS5, segments, events, profile), idempotent legacy-namespace migrations, plus path / namespace helpers (`sanitize_namespace`, `now_ts`, `namespace_dir`). -- **`documents.rs`** — `memory_docs` CRUD: `upsert_document` (chunks + embeds + writes markdown sidecar), `upsert_document_metadata_only` (light path), `list_documents`, `list_namespaces`, `delete_document`, `clear_namespace`. +- **`documents.rs`** — `memory_docs` CRUD: `upsert_document` (chunks + embeds + writes markdown sidecar), `upsert_documents` (the batch form: one embedding request per `EMBED_REQUEST_MAX_TEXTS` chunk texts across all the documents, then one transaction per document — tinymemory#138), `upsert_document_metadata_only` (light path), `list_documents`, `list_namespaces`, `delete_document`, `clear_namespace`. - **`kv.rs`** — global and namespace-scoped get/set/delete/list against `kv_global` / `kv_namespace`. - **`../../safety/`** — secret redaction/validation helpers. Document, KV, and episodic writes sanitize credentials before persistence and emit `[memory:safety]` diagnostics when a payload is rewritten. diff --git a/crates/tinymemory-core/src/store/namespace_store/documents.rs b/crates/tinymemory-core/src/store/namespace_store/documents.rs index 426c7e8c..80edaeda 100644 --- a/crates/tinymemory-core/src/store/namespace_store/documents.rs +++ b/crates/tinymemory-core/src/store/namespace_store/documents.rs @@ -1,8 +1,9 @@ //! Document CRUD against the `memory_docs` table. //! -//! Owns the upsert pipeline (with chunking + embedding), metadata-only writes -//! for high-frequency callers, list/delete/clear-namespace operations, and the -//! markdown sidecar files in `memory/namespaces//docs/`. +//! Owns the upsert pipeline (with chunking + embedding, batched across +//! documents), metadata-only writes for high-frequency callers, +//! list/delete/clear-namespace operations, and the markdown sidecar files in +//! `memory/namespaces//docs/`. use rusqlite::{params, OptionalExtension}; use serde_json::{json, Value}; @@ -14,11 +15,29 @@ use crate::store::types::{NamespaceDocumentInput, StoredMemoryDocument, GLOBAL_N use super::UnifiedMemory; +/// Token budget per vector chunk when a document is split for embedding. +const DOCUMENT_CHUNK_MAX_TOKENS: usize = 225; + +/// Upper bound on chunk texts sent to the embedding provider in one request +/// when a batch of documents is embedded together +/// ([`UnifiedMemory::upsert_documents_presanitized`]). +/// +/// Sized under every provider's per-request input cap (Cohere admits 96 texts, +/// the others more) and, at [`DOCUMENT_CHUNK_MAX_TOKENS`] per chunk, roughly +/// 14k tokens per request — under every provider's token budget — so a batch +/// of any length becomes a handful of bounded requests rather than one a +/// provider may refuse or time out. Still one to two orders of magnitude +/// fewer round-trips than the one-per-document write path it replaces. +pub(crate) const EMBED_REQUEST_MAX_TEXTS: usize = 64; + impl UnifiedMemory { /// Insert or update a document by `(namespace, key)`. Writes the markdown /// sidecar, replaces vector chunks, and embeds them with the configured /// provider. /// + /// The one-document case of [`Self::upsert_documents_presanitized`]; see it + /// for the write order and the failure contract. + /// /// **Takes already-sanitized input.** The host secret/PII write gate runs /// in [`crate::store::write_gate`], which owns this /// method's only call site; use `UnifiedMemory::upsert_document` instead @@ -27,6 +46,115 @@ impl UnifiedMemory { pub(crate) async fn upsert_document_presanitized( &self, input: NamespaceDocumentInput, + ) -> Result { + self.upsert_documents_presanitized(vec![input]) + .await + .pop() + .unwrap_or_else(|| Err("document upsert produced no result".to_string())) + } + + /// Insert or update many documents, embedding their chunks **together**: + /// one provider request per [`EMBED_REQUEST_MAX_TEXTS`] chunk texts across + /// the whole batch rather than one request per document (tinymemory#138). + /// A connector pass of several hundred small items used to pay one + /// embedding round-trip each; here it pays one per bounded group of texts. + /// + /// Order of operations: every document is chunked, every chunk text is + /// embedded (see [`Self::embed_chunk_texts`]), then the documents are + /// written one at a time in input order — each under its own per-key write + /// lock, with the row, the chunk replacement and the new vectors in ONE + /// transaction, so a reader never sees a row whose chunks are still being + /// replaced. + /// + /// The first document whose write fails ends the batch: the result holds + /// one entry per document attempted, in input order, with that failure as + /// its last entry, and the documents after it are left untouched. An + /// embedding failure is not a write failure — a request the provider + /// refuses leaves the chunks it covered vector-less (keyword-searchable and + /// re-embeddable), exactly as the single-document path always has, and the + /// batch carries on. + /// + /// **Takes already-sanitized input** — same contract as + /// [`Self::upsert_document_presanitized`]; go through + /// `UnifiedMemory::upsert_documents` instead. + pub(crate) async fn upsert_documents_presanitized( + &self, + inputs: Vec, + ) -> Vec> { + let chunked: Vec> = inputs + .iter() + .map(|input| Self::chunk_document_content(&input.content, DOCUMENT_CHUNK_MAX_TOKENS)) + .collect(); + let texts: Vec<&str> = chunked.iter().flatten().map(String::as_str).collect(); + let mut vectors = self.embed_chunk_texts(&texts).await.into_iter(); + + let mut results = Vec::with_capacity(inputs.len()); + for (input, chunks) in inputs.into_iter().zip(chunked) { + // `embed_chunk_texts` yields exactly one slot per text, in order, + // so the next `chunks.len()` slots are this document's. + let embeddings: Vec>> = vectors.by_ref().take(chunks.len()).collect(); + let result = self + .write_document_presanitized(input, chunks, embeddings) + .await; + let failed = result.is_err(); + results.push(result); + if failed { + break; + } + } + results + } + + /// Embed `texts` in requests of at most [`EMBED_REQUEST_MAX_TEXTS`], + /// returning one slot per input position. + /// + /// Failure handling keeps the per-chunk resilience the single-document + /// path always had: + /// * a request the provider refuses leaves every position it covered + /// `None` — logged, not propagated, because a vector-less chunk is + /// still keyword-searchable and re-embeddable while a failed write is + /// lost; + /// * a provider that returns fewer vectors than texts, or an empty vector + /// for a position (`NoopEmbedding`, NaN recovery), leaves those + /// positions `None` by position. + async fn embed_chunk_texts(&self, texts: &[&str]) -> Vec>> { + let mut out: Vec>> = Vec::with_capacity(texts.len()); + for request in texts.chunks(EMBED_REQUEST_MAX_TEXTS) { + log::debug!( + "[memory] batch-embedding {} chunk text(s) in one request", + request.len() + ); + match self.embedder.embed(request).await { + Ok(vectors) => { + let mut vectors = vectors + .into_iter() + .map(|vector| (!vector.is_empty()).then_some(vector)); + out.extend(request.iter().map(|_| vectors.next().flatten())); + } + Err(e) => { + log::warn!( + "[memory] batch embed failed for {} chunk text(s); storing them without vectors: {e}", + request.len() + ); + out.resize(out.len() + request.len(), None); + } + } + } + out + } + + /// Persist one document whose chunks and vectors were computed up front. + /// + /// Under the per-key write lock: resolve the document id and `created_at`, + /// write the markdown sidecar, then commit the `memory_docs` row, the chunk + /// replacement and the new `vector_chunks` rows in a single transaction. + /// `embeddings` is aligned to `chunks` by position; a missing or `None` + /// slot stores that chunk without a vector. + async fn write_document_presanitized( + &self, + input: NamespaceDocumentInput, + chunks: Vec, + mut embeddings: Vec>>, ) -> Result { let namespace = Self::sanitize_namespace(&input.namespace); // The logical (delimiter-preserving) namespace, PII-redacted the same @@ -43,16 +171,17 @@ impl UnifiedMemory { if key.is_empty() { return Err("document key cannot be empty".to_string()); } - // Serialise writers of one key for the WHOLE operation. A deterministic + // Serialise writers of one key for the WHOLE write. A deterministic // document id stops two writers orphaning each other's chunks, but it - // does not ORDER them: the row write and the chunk replacement are - // separated by embedding, which awaits. Without this, writer A can - // update the row, await the embedder, and have B update the row and - // replace the chunks in between -- leaving B's content beside A's - // chunks, plus A's trailing chunks if A had more. The metadata-only - // path below takes the same lock: it writes the same row, so it must - // not interleave with a full write either. Same guard shape as the - // sync path's per-connection lock. + // does not ORDER them: the id / `created_at` lookups, the sidecar write + // (which awaits) and the transaction below must not interleave with + // another writer of the same key, or writer B's row can land between + // writer A's lookups and A's commit and be overwritten by content A + // resolved against stale state. The metadata-only path below takes the + // same lock: it writes the same row, so it must not interleave with a + // full write either. Same guard shape as the sync path's per-connection + // lock. Embedding happens before this lock is taken, so no provider + // round-trip is ever awaited while holding it. let _write_guard = Self::document_write_lock(&self.db_path, &namespace, &key) .lock_owned() .await; @@ -113,6 +242,8 @@ impl UnifiedMemory { let tags_json = serde_json::to_string(&input.tags).map_err(|e| e.to_string())?; let metadata_json = input.metadata.to_string(); + // Computed once; only attached to chunks that actually got a vector. + let signature = self.embedder.signature(); { let conn = self.conn.lock(); let tx = conn @@ -161,77 +292,36 @@ impl UnifiedMemory { params![namespace, document_id], ) .map_err(|e| format!("clear vector chunks: {e}"))?; - tx.commit().map_err(|e| format!("commit tx: {e}"))?; - } - - let chunks = Self::chunk_document_content(&input.content, 225); - - // Embed every chunk in a SINGLE provider call rather than one - // round-trip per chunk. All providers implement the batch `embed` - // (`embed_one` is just a convenience wrapper around it), so a document - // that chunks into N pieces previously paid N sequential network - // round-trips on the write path; this collapses them to one. - // - // Result handling preserves the previous per-chunk resilience: - // * a failed batch (provider error) stores all chunks WITHOUT a - // vector — exactly what `embed_one(...).await.ok()` did per chunk; - // * a provider that returns fewer/empty vectors than chunks (e.g. - // `NoopEmbedding` returns an empty Vec, or a blank position from - // NaN recovery) leaves those chunks vector-less by position. - let mut embeddings: Vec>> = if chunks.is_empty() { - Vec::new() - } else { - let chunk_refs: Vec<&str> = chunks.iter().map(String::as_str).collect(); - log::debug!( - "[memory] batch-embedding {} chunk(s) for {namespace}/{document_id}", - chunk_refs.len() - ); - match self.embedder.embed(&chunk_refs).await { - Ok(vectors) => vectors - .into_iter() - .map(|v| (!v.is_empty()).then_some(v)) - .collect(), - Err(e) => { - log::warn!( - "[memory] batch embed failed for {} chunk(s) in {namespace}/{document_id}; storing without vectors: {e}", - chunks.len() - ); - Vec::new() - } + for (idx, chunk) in chunks.iter().enumerate() { + // Move the vector out by position so recall can exclude vectors + // produced by a different embedding model (cross-model cosine is + // meaningless) and guard against dimension mismatches. Missing + // positions (short/empty provider result) stay vector-less. + let embedded = embeddings.get_mut(idx).and_then(Option::take); + let dim = embedded.as_ref().map(|v| v.len() as i64); + let model_signature = embedded.as_ref().map(|_| signature.clone()); + let embedding = embedded.as_ref().map(|v| Self::vec_to_bytes(v)); + let chunk_id = format!("{document_id}:{idx}"); + tx.execute( + "INSERT OR REPLACE INTO vector_chunks + (namespace, document_id, chunk_id, text, embedding, metadata_json, created_at, updated_at, model_signature, dim) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)", + params![ + namespace, + document_id, + chunk_id, + chunk, + embedding, + json!({"lancedb_table": format!("ns_{namespace}"), "chunk_index": idx}).to_string(), + now, + now, + model_signature, + dim + ], + ) + .map_err(|e| format!("insert vector chunk: {e}"))?; } - }; - - // Computed once; only attached to chunks that actually got a vector. - let signature = self.embedder.signature(); - for (idx, chunk) in chunks.iter().enumerate() { - // Move the vector out by position so recall can exclude vectors - // produced by a different embedding model (cross-model cosine is - // meaningless) and guard against dimension mismatches. Missing - // positions (short/empty provider result) stay vector-less. - let embedded = embeddings.get_mut(idx).and_then(Option::take); - let dim = embedded.as_ref().map(|v| v.len() as i64); - let model_signature = embedded.as_ref().map(|_| signature.clone()); - let embedding = embedded.as_ref().map(|v| Self::vec_to_bytes(v)); - let chunk_id = format!("{document_id}:{idx}"); - let conn = self.conn.lock(); - conn.execute( - "INSERT OR REPLACE INTO vector_chunks - (namespace, document_id, chunk_id, text, embedding, metadata_json, created_at, updated_at, model_signature, dim) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)", - params![ - namespace, - document_id, - chunk_id, - chunk, - embedding, - json!({"lancedb_table": format!("ns_{namespace}"), "chunk_index": idx}).to_string(), - now, - now, - model_signature, - dim - ], - ) - .map_err(|e| format!("insert vector chunk: {e}"))?; + tx.commit().map_err(|e| format!("commit tx: {e}"))?; } Ok(document_id) diff --git a/crates/tinymemory-core/src/store/namespace_store/documents_tests.rs b/crates/tinymemory-core/src/store/namespace_store/documents_tests.rs index 7716c69e..f965d3fc 100644 --- a/crates/tinymemory-core/src/store/namespace_store/documents_tests.rs +++ b/crates/tinymemory-core/src/store/namespace_store/documents_tests.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use serde_json::json; use tempfile::TempDir; +use super::EMBED_REQUEST_MAX_TEXTS; use crate::store::{NamespaceDocumentInput, UnifiedMemory}; use tinymemory_api::host::NoopEmbedding; @@ -40,6 +41,18 @@ fn count_vector_chunks(memory: &UnifiedMemory, namespace: &str, document_id: &st .unwrap() } +/// Like [`count_vector_chunks`], counting only the chunks that carry a vector. +fn count_embedded_chunks(memory: &UnifiedMemory, namespace: &str, document_id: &str) -> i64 { + let conn = memory.conn.lock(); + conn.query_row( + "SELECT COUNT(*) FROM vector_chunks + WHERE namespace = ?1 AND document_id = ?2 AND embedding IS NOT NULL", + rusqlite::params![UnifiedMemory::sanitize_namespace(namespace), document_id], + |row| row.get(0), + ) + .unwrap() +} + #[tokio::test] async fn list_documents_without_namespace_returns_all_docs_across_namespaces() { let tmp = TempDir::new().unwrap(); @@ -421,6 +434,196 @@ async fn upsert_document_batch_embeds_all_chunks_in_one_call() { ); } +/// Embedder that records the size of every request and can refuse one of them +/// (by 0-based request index), so the batch path's request splitting and its +/// per-request failure isolation are observable. +struct RequestRecordingEmbedder { + requests: std::sync::Mutex>, + fail_request: Option, +} + +impl RequestRecordingEmbedder { + fn new(fail_request: Option) -> Arc { + Arc::new(Self { + requests: std::sync::Mutex::new(Vec::new()), + fail_request, + }) + } + + fn requests(&self) -> Vec { + self.requests.lock().unwrap().clone() + } +} + +#[async_trait::async_trait] +impl tinymemory_api::host::EmbeddingProvider for RequestRecordingEmbedder { + fn name(&self) -> &str { + "recording" + } + + fn model_id(&self) -> &str { + "recording-test" + } + + fn dimensions(&self) -> usize { + 3 + } + + async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { + let index = { + let mut requests = self.requests.lock().unwrap(); + requests.push(texts.len()); + requests.len() - 1 + }; + if self.fail_request == Some(index) { + anyhow::bail!("provider refused request {index}"); + } + Ok(texts.iter().map(|_| vec![0.1, 0.2, 0.3]).collect()) + } +} + +/// tinymemory#138: a batch of documents must share embedding requests rather +/// than pay one round-trip per document. +#[tokio::test] +async fn upsert_documents_embeds_every_document_in_one_request() { + let tmp = TempDir::new().unwrap(); + let embedder = RequestRecordingEmbedder::new(None); + let memory = UnifiedMemory::new(tmp.path(), embedder.clone(), None).unwrap(); + + // Each body chunks into several pieces, so the batch has many more chunks + // than documents — and still fits one request. + let long_body = "alpha ".repeat(400); + let inputs: Vec = ["doc-a", "doc-b", "doc-c"] + .iter() + .map(|key| make_doc_input("test:batch-many", key, key, &long_body)) + .collect(); + + let ids: Vec = memory + .upsert_documents(inputs) + .await + .into_iter() + .map(|result| result.unwrap()) + .collect(); + assert_eq!(ids.len(), 3); + + let mut total_chunks = 0; + for id in &ids { + let chunks = count_vector_chunks(&memory, "test:batch-many", id); + assert!( + chunks >= 3, + "each body should chunk into >=3 pieces, got {chunks}" + ); + assert_eq!( + count_embedded_chunks(&memory, "test:batch-many", id), + chunks, + "every chunk of every document must carry a vector" + ); + total_chunks += chunks; + } + assert_eq!( + embedder.requests(), + vec![total_chunks as usize], + "three documents' chunks must travel in ONE provider request, not one per document" + ); +} + +#[tokio::test] +async fn upsert_documents_splits_embedding_requests_at_the_request_cap() { + let tmp = TempDir::new().unwrap(); + let embedder = RequestRecordingEmbedder::new(None); + let memory = UnifiedMemory::new(tmp.path(), embedder.clone(), None).unwrap(); + + // One chunk per document, one more document than a request may carry. + let inputs: Vec = (0..=EMBED_REQUEST_MAX_TEXTS) + .map(|n| make_doc_input("test:cap", &format!("doc-{n}"), "Doc", &format!("body {n}"))) + .collect(); + + let results = memory.upsert_documents(inputs).await; + assert_eq!(results.len(), EMBED_REQUEST_MAX_TEXTS + 1); + assert!(results.iter().all(Result::is_ok), "{results:?}"); + assert_eq!( + embedder.requests(), + vec![EMBED_REQUEST_MAX_TEXTS, 1], + "chunk texts must be sent in requests of at most EMBED_REQUEST_MAX_TEXTS" + ); +} + +#[tokio::test] +async fn upsert_documents_keeps_writing_when_one_embedding_request_is_refused() { + let tmp = TempDir::new().unwrap(); + // The second request (the lone overflow chunk) is refused. + let embedder = RequestRecordingEmbedder::new(Some(1)); + let memory = UnifiedMemory::new(tmp.path(), embedder.clone(), None).unwrap(); + + let inputs: Vec = (0..=EMBED_REQUEST_MAX_TEXTS) + .map(|n| { + make_doc_input( + "test:refused", + &format!("doc-{n}"), + "Doc", + &format!("body {n}"), + ) + }) + .collect(); + + let ids: Vec = memory + .upsert_documents(inputs) + .await + .into_iter() + .map(|result| result.expect("an embedding failure is not a write failure")) + .collect(); + assert_eq!(ids.len(), EMBED_REQUEST_MAX_TEXTS + 1); + assert_eq!(embedder.requests(), vec![EMBED_REQUEST_MAX_TEXTS, 1]); + + let first = &ids[0]; + assert_eq!(count_vector_chunks(&memory, "test:refused", first), 1); + assert_eq!( + count_embedded_chunks(&memory, "test:refused", first), + 1, + "documents covered by the accepted request keep their vectors" + ); + let last = ids.last().unwrap(); + assert_eq!( + count_vector_chunks(&memory, "test:refused", last), + 1, + "the document covered by the refused request is still written" + ); + assert_eq!( + count_embedded_chunks(&memory, "test:refused", last), + 0, + "…but its chunk is stored without a vector, like the single-document path" + ); +} + +#[tokio::test] +async fn upsert_documents_stops_at_the_first_write_failure() { + let tmp = TempDir::new().unwrap(); + let memory = UnifiedMemory::new(tmp.path(), Arc::new(NoopEmbedding), None).unwrap(); + + let results = memory + .upsert_documents(vec![ + make_doc_input("test:stop", "doc-a", "Doc A", "A body"), + make_doc_input("test:stop", " ", "Blank key", "rejected by the store"), + make_doc_input("test:stop", "doc-c", "Doc C", "never attempted"), + ]) + .await; + + assert_eq!( + results.len(), + 2, + "the failing document is the last entry; nothing after it is attempted" + ); + assert!(results[0].is_ok()); + let err = results[1].as_ref().unwrap_err(); + assert!( + err.contains("document key cannot be empty"), + "the failing entry carries the store's own error, got {err:?}" + ); + let docs = memory.list_documents(Some("test:stop")).await.unwrap(); + assert_eq!(docs["count"].as_u64(), Some(1)); + assert_eq!(docs["documents"][0]["key"], "doc-a"); +} + #[tokio::test] async fn upsert_document_reuses_document_id_preserves_created_at_and_replaces_vector_chunks() { let tmp = TempDir::new().unwrap(); diff --git a/crates/tinymemory-core/src/store/write_gate.rs b/crates/tinymemory-core/src/store/write_gate.rs index 1f204c36..1f04ed25 100644 --- a/crates/tinymemory-core/src/store/write_gate.rs +++ b/crates/tinymemory-core/src/store/write_gate.rs @@ -18,7 +18,8 @@ //! `upsert_document_metadata_only` as inherent methods on `UnifiedMemory` with //! the **same names and signatures they always had**, so every existing caller //! is routed through the gate without a single call-site edit — which is also -//! what makes the "no bypass" claim below checkable rather than hopeful. +//! what makes the "no bypass" claim below checkable rather than hopeful. The +//! batch form, `upsert_documents`, is declared here for the same reason. //! //! # The gate, in order //! @@ -45,12 +46,14 @@ //! //! # No bypass //! -//! `upsert_document_presanitized` / `upsert_document_metadata_only_presanitized` -//! are `pub(crate)` and newly named, and this module holds their only call -//! sites — verify with: +//! `upsert_document_presanitized` / `upsert_documents_presanitized` / +//! `upsert_document_metadata_only_presanitized` are `pub(crate)` and newly +//! named, and this module holds their only call sites outside `documents.rs` +//! itself (where the one-document method is the one-element case of the batch +//! one) — verify with: //! //! ```text -//! rg 'upsert_document(_metadata_only)?_presanitized' src/ +//! rg 'upsert_documents?(_metadata_only)?_presanitized' src/ //! ``` //! //! Every other writer in the tree (`Memory::store_with_taint`, `MemoryClient`, @@ -142,6 +145,48 @@ impl UnifiedMemory { } } + /// Insert or update many documents, applying the host secret/PII write + /// gate to each and embedding their chunks together — one provider request + /// per bounded group of chunk texts across the batch rather than one per + /// document (tinymemory#138). + /// + /// Inputs are gated in order and the admitted prefix is written by + /// `Self::upsert_documents_presanitized`, which stops at the first write + /// failure; the first gate rejection, if any, ends the batch the same way. + /// The result holds one entry per document attempted, in input order, so a + /// failure is always the last entry and every document before it was + /// written. + /// + /// # Errors + /// + /// Per document, the same failure modes as [`Self::upsert_document`]. + pub async fn upsert_documents( + &self, + inputs: Vec, + ) -> Vec> { + let mut admitted = Vec::with_capacity(inputs.len()); + let mut rejection = None; + for input in inputs { + match gate(input, "document") { + GateOutcome::Admit(input) => admitted.push(*input), + GateOutcome::Reject(err) => { + rejection = Some(err); + break; + } + } + } + let mut results = self.upsert_documents_presanitized(admitted).await; + if let Some(err) = rejection { + // Only when the admitted prefix was written in full: a write failure + // in it is already the batch's last entry, and everything after a + // failure — the rejected document included — stays unattempted. + if results.iter().all(Result::is_ok) { + results.push(Err(err)); + } + } + results + } + /// Store a document without chunking, embedding, or graph extraction, /// applying the host secret/PII write gate first. /// diff --git a/crates/tinymemory-core/src/store/write_gate_tests.rs b/crates/tinymemory-core/src/store/write_gate_tests.rs index 9048d6f1..a1789695 100644 --- a/crates/tinymemory-core/src/store/write_gate_tests.rs +++ b/crates/tinymemory-core/src/store/write_gate_tests.rs @@ -174,3 +174,39 @@ async fn gate_canonicalizes_a_pii_like_key_and_keeps_the_row_addressable() { docs[0].key ); } + +#[tokio::test] +async fn gated_batch_upsert_redacts_each_document_and_stops_at_the_first_rejection() { + let (_tmp, memory) = fresh(); + + let mut secret_key = secret_doc("sk-1234567890123456789012345"); + secret_key.namespace = "safe".to_string(); + let results = memory + .upsert_documents(vec![ + secret_doc("first"), + secret_key, + secret_doc("never-attempted"), + ]) + .await; + + assert_eq!( + results.len(), + 2, + "the rejected document ends the batch as its last entry, got {results:?}" + ); + assert!(results[0].is_ok(), "{results:?}"); + let err = results[1].as_ref().unwrap_err(); + assert!( + err.contains("cannot contain secrets"), + "secret-like key must be refused, got {err:?}" + ); + + let docs = memory.load_documents_for_scope("safe").await.unwrap(); + assert_eq!(docs.len(), 1, "only the admitted prefix reaches storage"); + assert_eq!(docs[0].key, "first"); + assert!( + !docs[0].content.contains("BEGIN PRIVATE KEY"), + "a batched write must be redacted exactly like a single one, got {:?}", + docs[0].content + ); +} diff --git a/crates/tinymemory-module/Cargo.lock b/crates/tinymemory-module/Cargo.lock index c4f8de80..bc426e5a 100644 --- a/crates/tinymemory-module/Cargo.lock +++ b/crates/tinymemory-module/Cargo.lock @@ -266,34 +266,13 @@ dependencies = [ "crypto-common 0.2.2", ] -[[package]] -name = "dirs" -version = "5.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44c45a9d03d6676652bcb5e724c7e988de1acad23a711b5217ab9cbecbec2225" -dependencies = [ - "dirs-sys 0.4.1", -] - [[package]] name = "dirs" version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3e8aa94d75141228480295a7d0e7feb620b1a5ad9f12bc40be62411e38cce4e" dependencies = [ - "dirs-sys 0.5.0", -] - -[[package]] -name = "dirs-sys" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "520f05a5cbd335fae5a99ff7a6ab8627577660ee5cfd6a94a6a929b52ff0321c" -dependencies = [ - "libc", - "option-ext", - "redox_users 0.4.6", - "windows-sys 0.48.0", + "dirs-sys", ] [[package]] @@ -304,7 +283,7 @@ checksum = "e01a3366d27ee9890022452ee61b2b63a67e6f13f58900b651ff5665f0bb1fab" dependencies = [ "libc", "option-ext", - "redox_users 0.5.2", + "redox_users", "windows-sys 0.61.2", ] @@ -1117,7 +1096,7 @@ dependencies = [ "rustc-hash", "rustls", "socket2", - "thiserror 2.0.20", + "thiserror", "tokio", "tracing", "web-time", @@ -1139,7 +1118,7 @@ dependencies = [ "rustls", "rustls-pki-types", "slab", - "thiserror 2.0.20", + "thiserror", "tinyvec", "tracing", "web-time", @@ -1209,17 +1188,6 @@ dependencies = [ "bitflags", ] -[[package]] -name = "redox_users" -version = "0.4.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba009ff324d1fc1b900bd1fdb31564febe58a8ccc8a6fdbb93b543d33b13ca43" -dependencies = [ - "getrandom 0.2.17", - "libredox", - "thiserror 1.0.69", -] - [[package]] name = "redox_users" version = "0.5.2" @@ -1228,7 +1196,7 @@ checksum = "a4e608c6638b9c18977b00b475ac1f28d14e84b27d8d42f70e0bf1e3dec127ac" dependencies = [ "getrandom 0.2.17", "libredox", - "thiserror 2.0.20", + "thiserror", ] [[package]] @@ -1342,7 +1310,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c51c9ae4df8a7fba42103df5c621fa3c37eccf3a3c650879e90fc48b11cc192c" dependencies = [ "hashbrown 0.16.1", - "thiserror 2.0.20", + "thiserror", ] [[package]] @@ -1707,33 +1675,13 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "thiserror" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" -dependencies = [ - "thiserror-impl 1.0.69", -] - [[package]] name = "thiserror" version = "2.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" dependencies = [ - "thiserror-impl 2.0.20", -] - -[[package]] -name = "thiserror-impl" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.119", + "thiserror-impl", ] [[package]] @@ -1757,7 +1705,7 @@ dependencies = [ "serde_json", "tar", "tempfile", - "thiserror 2.0.20", + "thiserror", "tinybus-macros", "tokio", "toml 0.8.23", @@ -1789,13 +1737,13 @@ dependencies = [ [[package]] name = "tinycortex" -version = "0.1.1" +version = "0.1.2" dependencies = [ "anyhow", "async-trait", "block2", "chrono", - "dirs 5.0.1", + "dirs", "futures", "git2", "hex", @@ -1812,7 +1760,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.10.9", - "thiserror 2.0.20", + "thiserror", "tinycortex-api", "tinyinference", "tokio", @@ -1841,7 +1789,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.11.0", - "thiserror 2.0.20", + "thiserror", "tokio", ] @@ -1879,7 +1827,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.11.0", - "thiserror 2.0.20", + "thiserror", "uuid", ] @@ -1890,7 +1838,7 @@ dependencies = [ "anyhow", "async-trait", "chrono", - "dirs 6.0.0", + "dirs", "futures", "log", "parking_lot", @@ -1901,7 +1849,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.11.0", - "thiserror 2.0.20", + "thiserror", "tinycortex", "tinycortex-api", "tinyinference", @@ -2450,7 +2398,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] @@ -2512,22 +2460,13 @@ dependencies = [ "windows-link", ] -[[package]] -name = "windows-sys" -version = "0.48.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9" -dependencies = [ - "windows-targets 0.48.5", -] - [[package]] name = "windows-sys" version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets 0.52.6", + "windows-targets", ] [[package]] @@ -2539,67 +2478,34 @@ dependencies = [ "windows-link", ] -[[package]] -name = "windows-targets" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" -dependencies = [ - "windows_aarch64_gnullvm 0.48.5", - "windows_aarch64_msvc 0.48.5", - "windows_i686_gnu 0.48.5", - "windows_i686_msvc 0.48.5", - "windows_x86_64_gnu 0.48.5", - "windows_x86_64_gnullvm 0.48.5", - "windows_x86_64_msvc 0.48.5", -] - [[package]] name = "windows-targets" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ - "windows_aarch64_gnullvm 0.52.6", - "windows_aarch64_msvc 0.52.6", - "windows_i686_gnu 0.52.6", + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", "windows_i686_gnullvm", - "windows_i686_msvc 0.52.6", - "windows_x86_64_gnu 0.52.6", - "windows_x86_64_gnullvm 0.52.6", - "windows_x86_64_msvc 0.52.6", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", ] -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" - [[package]] name = "windows_aarch64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" -[[package]] -name = "windows_aarch64_msvc" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" - [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" -[[package]] -name = "windows_i686_gnu" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" - [[package]] name = "windows_i686_gnu" version = "0.52.6" @@ -2612,48 +2518,24 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" -[[package]] -name = "windows_i686_msvc" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" - [[package]] name = "windows_i686_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" -[[package]] -name = "windows_x86_64_gnu" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" - [[package]] name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" - [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" -[[package]] -name = "windows_x86_64_msvc" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" - [[package]] name = "windows_x86_64_msvc" version = "0.52.6" @@ -2787,7 +2669,7 @@ dependencies = [ "flate2", "indexmap", "memchr", - "thiserror 2.0.20", + "thiserror", "zopfli", ] diff --git a/crates/tinymemory-tinycortex/src/engine/mod.rs b/crates/tinymemory-tinycortex/src/engine/mod.rs index de4a0b3d..c32f80e8 100644 --- a/crates/tinymemory-tinycortex/src/engine/mod.rs +++ b/crates/tinymemory-tinycortex/src/engine/mod.rs @@ -2198,11 +2198,25 @@ impl MemorySourceSink for TinycortexProvider { // Resolved once: whether a source's items belong in the memory tree is a // property of the source, not of the item. let tree_scope = connector_tree_scope(source_kind, source_id); + + // Convert every item up front so the store can embed the whole batch + // together (tinymemory#138): `put_docs` pays one embedding round-trip + // per bounded group of chunk texts across ALL the items instead of one + // per item, which is what put a 500-item connector pass at ~15 minutes + // and over the host's slow-call deadline. An item that cannot be + // converted (a blank id, a shape the store rejects) ends the conversion + // where it stands: the items before it are still written and treed + // below, then its error is returned — the same end state as when each + // item was written as it was reached. + let mut inputs = Vec::with_capacity(items_len); + let mut tree_items = Vec::with_capacity(items_len); + let mut rejected = None; for item in items { if item.item_id.trim().is_empty() { - return Err(MemoryError::Invalid( + rejected = Some(MemoryError::Invalid( "source item_id must not be empty".to_string(), )); + break; } let title = if item.title.trim().is_empty() { item.item_id.clone() @@ -2235,8 +2249,38 @@ impl MemorySourceSink for TinycortexProvider { document_id: None, taint, }; - let input = Self::cross(&input, "convert source document")?; - match self.client.put_doc(input).await { + match Self::cross(&input, "convert source document") { + Ok(input) => { + inputs.push(input); + tree_items.push(tree_item); + } + Err(error) => { + rejected = Some(error); + break; + } + } + } + + // One store call for the whole batch. Its results hold one entry per + // document attempted, in order, with a failed write always last, so + // walking them in order keeps exactly the per-item accounting the old + // one-write-per-item loop had. + let batch = self.client.put_docs(inputs).await; + if batch.dropped_extractions > 0 { + // Best-effort by the queue's contract: the documents and (below) + // their memory-tree chunks are stored; only the namespace graph + // extraction for these items was skipped. Named per source so an + // operator can tell which sync overran the queue. + log::warn!( + "[tinycortex:sources] graph extraction skipped for {} of {} written \ + item(s) of source `{source_id}`: the ingestion queue refused the \ + job(s); the documents and their memory-tree chunks are stored", + batch.dropped_extractions, + batch.results.iter().filter(|result| result.is_ok()).count() + ); + } + for (result, tree_item) in batch.results.into_iter().zip(tree_items) { + match result { Ok(id) => { outcome.written = outcome.written.saturating_add(1); outcome.ids.push(id); @@ -2296,7 +2340,10 @@ impl MemorySourceSink for TinycortexProvider { } } } - Ok(outcome) + match rejected { + Some(error) => Err(error), + None => Ok(outcome), + } } async fn forget_source(&self, source_id: &str) -> Result { diff --git a/crates/tinymemory-tinycortex/tests/full_provider_conformance.rs b/crates/tinymemory-tinycortex/tests/full_provider_conformance.rs index 3b8e8258..1e6b239c 100644 --- a/crates/tinymemory-tinycortex/tests/full_provider_conformance.rs +++ b/crates/tinymemory-tinycortex/tests/full_provider_conformance.rs @@ -4016,3 +4016,154 @@ async fn a_connector_sync_reaches_the_memory_tree_and_is_forgotten_with_its_sour (openhuman#6007)" ); } + +/// Embedder that records the size of every request, so a test can prove how +/// many provider round-trips a batch of source items paid for. +struct RequestCountingEmbedder { + requests: std::sync::Mutex>, +} + +#[async_trait::async_trait] +impl tinymemory_api::host::EmbeddingProvider for RequestCountingEmbedder { + fn name(&self) -> &str { + "counting" + } + + fn model_id(&self) -> &str { + "counting-test" + } + + fn dimensions(&self) -> usize { + 3 + } + + async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { + self.requests + .lock() + .expect("requests lock") + .push(texts.len()); + Ok(texts.iter().map(|_| vec![0.1, 0.2, 0.3]).collect()) + } +} + +/// A provider whose document store embeds through `embedder`. The store takes +/// its embedder directly, so this bypasses the process-global seam and leaves +/// what every other test in this binary sees untouched. +fn provider_with_embedder( + workspace: &std::path::Path, + embedder: Arc, +) -> TinycortexProvider { + tinymemory_core::embedding_host::set_embedding_host(Arc::new(NoopEmbeddingHost)); + let memory = tinymemory_core::store::UnifiedMemory::new(workspace, embedder, None) + .expect("open the workspace store"); + let client = Arc::new(tinymemory_core::store::MemoryClient::from_unified_memory( + memory, + )); + TinycortexProvider::new( + "tinycortex".into(), + provider_config(workspace, serde_json::Value::Null), + client, + ) +} + +/// tinymemory#138: a connector pass hands the sink hundreds of small items, and +/// each one used to pay its own embedding round-trip — about 1.7 s per item +/// against the managed embedder, so a 500-item pass ran into the host's +/// 15-minute slow-call deadline. The batch must share requests across items. +#[tokio::test(flavor = "multi_thread")] +async fn source_items_are_embedded_together_rather_than_one_request_per_item() { + use tinymemory_api::provider::types::SourceItem; + use tinymemory_api::provider::{ChunkQuery, MemoryProvider}; + use tinymemory_api::types::MemoryTaint; + + let workspace = tempfile::tempdir().expect("workspace"); + let embedder = Arc::new(RequestCountingEmbedder { + requests: std::sync::Mutex::new(Vec::new()), + }); + let provider = provider_with_embedder(workspace.path(), Arc::clone(&embedder)); + let source = provider.as_sources().expect("SourceSink"); + + let items: Vec = (1..=5) + .map(|n| SourceItem { + item_id: format!("msg-{n}"), + title: format!("Message {n}"), + content: format!("Short message number {n} about the roadmap."), + mime: Some("text/plain".into()), + url: None, + updated_at_ms: Some(n), + tags: vec!["gmail".into()], + }) + .collect(); + let outcome = source + .accept_source_items("gmail:conn-1", "composio", items, MemoryTaint::ExternalSync) + .await + .expect("accept the batch"); + assert_eq!(outcome.written, 5); + assert_eq!(outcome.ids.len(), 5, "one id per written item, in order"); + + let requests = embedder.requests.lock().expect("requests lock").clone(); + assert_eq!( + requests, + vec![5], + "five one-chunk items must cost ONE embedding request, not five (tinymemory#138); \ + got {requests:?}" + ); + // The tree funnel still runs once per written item (openhuman#6007). + let chunks = provider.as_chunks().expect("Chunks"); + assert_eq!( + chunks + .count_chunks(&ChunkQuery::default(), None) + .await + .expect("count chunks"), + 5, + "every written item must still reach the memory tree" + ); +} + +/// The batch keeps the sink's per-item accounting: an item the sink rejects +/// fails the call, after the items before it were written and none after it. +#[tokio::test(flavor = "multi_thread")] +async fn a_blank_source_item_id_fails_the_batch_after_the_items_before_it() { + use tinymemory_api::error::MemoryError; + use tinymemory_api::provider::types::SourceItem; + use tinymemory_api::provider::MemoryProvider; + use tinymemory_api::types::MemoryTaint; + + let workspace = tempfile::tempdir().expect("workspace"); + let provider = provider_over(workspace.path()); + let source = provider.as_sources().expect("SourceSink"); + let item = |id: &str| SourceItem { + item_id: id.into(), + title: "Item".into(), + content: "body".into(), + mime: None, + url: None, + updated_at_ms: None, + tags: Vec::new(), + }; + + let error = source + .accept_source_items( + "drive-1", + "drive", + vec![item("a"), item("b"), item(" "), item("d")], + MemoryTaint::ExternalSync, + ) + .await + .expect_err("a blank item id must fail the call"); + assert!( + matches!(&error, MemoryError::Invalid(message) if message.contains("item_id must not be empty")), + "got {error:?}" + ); + + let documents = provider.as_documents().expect("Documents"); + let listed = documents + .list_documents(Some("source:drive-1")) + .await + .expect("list documents"); + assert_eq!( + listed["count"].as_u64(), + Some(2), + "the items before the blank one are written; the ones after it are not" + ); +} diff --git a/vendor/tinycortex b/vendor/tinycortex index cb1c1639..79131f27 160000 --- a/vendor/tinycortex +++ b/vendor/tinycortex @@ -1 +1 @@ -Subproject commit cb1c1639c8a82221e1d8e8a9715b7640bb3f05d8 +Subproject commit 79131f275cf98f06421fefd9dc812c5a63a5a182