From 6f1721cb598afd0ae4c03408b263ca8203a19b87 Mon Sep 17 00:00:00 2001 From: mysma-9403 <64923976+mysma-9403@users.noreply.github.com> Date: Wed, 3 Jun 2026 00:31:45 +0200 Subject: [PATCH] perf(memory_store): batch chunk embeddings in upsert_document (#3213) --- .../memory_store/unified/documents.rs | 49 +++++++++++++++-- .../memory_store/unified/documents_tests.rs | 54 +++++++++++++++++++ 2 files changed, 98 insertions(+), 5 deletions(-) diff --git a/src/openhuman/memory_store/unified/documents.rs b/src/openhuman/memory_store/unified/documents.rs index 38f5226a0..7e212e4fc 100644 --- a/src/openhuman/memory_store/unified/documents.rs +++ b/src/openhuman/memory_store/unified/documents.rs @@ -154,13 +154,52 @@ impl UnifiedMemory { } 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() + } + } + }; + + // Computed once; only attached to chunks that actually got a vector. + let signature = self.embedder.signature(); for (idx, chunk) in chunks.iter().enumerate() { - // Embed the chunk, capturing the model signature + dimension so recall - // can exclude vectors produced by a different embedding model (cross-model - // cosine is meaningless) and guard against dimension mismatches. - let embedded = self.embedder.embed_one(chunk).await.ok(); + // 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(|_| self.embedder.signature()); + 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(); diff --git a/src/openhuman/memory_store/unified/documents_tests.rs b/src/openhuman/memory_store/unified/documents_tests.rs index 81da179f2..6e9936ad5 100644 --- a/src/openhuman/memory_store/unified/documents_tests.rs +++ b/src/openhuman/memory_store/unified/documents_tests.rs @@ -366,6 +366,60 @@ async fn upsert_document_writes_vector_chunks_for_chunked_content() { ); } +/// Embedder that records how many times `embed` is invoked and returns one +/// fixed-dimension vector per input text. Used to prove `upsert_document` +/// embeds all chunks in a SINGLE batch call rather than one call per chunk. +struct CountingEmbedder { + calls: std::sync::atomic::AtomicUsize, +} + +#[async_trait::async_trait] +impl crate::openhuman::embeddings::EmbeddingProvider for CountingEmbedder { + 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.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(texts.iter().map(|_| vec![0.1, 0.2, 0.3]).collect()) + } +} + +#[tokio::test] +async fn upsert_document_batch_embeds_all_chunks_in_one_call() { + let tmp = TempDir::new().unwrap(); + let embedder = Arc::new(CountingEmbedder { + calls: std::sync::atomic::AtomicUsize::new(0), + }); + let memory = UnifiedMemory::new(tmp.path(), embedder.clone(), None).unwrap(); + + // Long enough to chunk into several pieces (chunk size is 225 chars). + let long_body = "alpha ".repeat(400); + let document_id = memory + .upsert_document(make_doc_input("test:batch", "doc-a", "Doc A", &long_body)) + .await + .unwrap(); + + let chunk_count = count_vector_chunks(&memory, "test:batch", &document_id); + assert!( + chunk_count >= 3, + "test body should chunk into >=3 pieces, got {chunk_count}" + ); + assert_eq!( + embedder.calls.load(std::sync::atomic::Ordering::SeqCst), + 1, + "all chunks must be embedded in a single batch call, not one call per chunk" + ); +} + #[tokio::test] async fn upsert_document_reuses_document_id_preserves_created_at_and_replaces_vector_chunks() { let tmp = TempDir::new().unwrap();