perf(memory_store): batch chunk embeddings in upsert_document (#3213)

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<Option<Vec<f32>>> = 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();

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<Vec<Vec<f32>>> {
+        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();