feat(memory): defer embeddings to post-sync batch pass (#3428)

src/core/jsonrpc.rs

@@ -1892,7 +1892,7 @@ fn register_domain_subscribers(
         crate::openhuman::memory_conversations::register_conversation_persistence_subscriber(
             workspace_dir.clone(),
         );
-        crate::openhuman::memory::sync::register_sync_stage_bridge();
+        crate::openhuman::memory::sync::register_sync_stage_bridge(&config);
         if let Err(error) = crate::openhuman::composio::init_composio_trigger_history(
             workspace_dir.clone(),
         ) {

src/openhuman/channels/runtime/startup.rs

@@ -87,7 +87,7 @@ pub async fn start_channels(mut config: Config) -> Result<()> {
     crate::openhuman::memory_conversations::register_conversation_persistence_subscriber(
         config.workspace_dir.clone(),
     );
-    crate::openhuman::memory::sync::register_sync_stage_bridge();
+    crate::openhuman::memory::sync::register_sync_stage_bridge(&config);
     crate::openhuman::composio::register_composio_trigger_subscriber();
     // Surface parked ApprovalGate requests as chat messages so the user can
     // answer yes/no in the thread (chat-native approval, issue #1339).

src/openhuman/memory/sync.rs

@@ -124,10 +124,12 @@ pub fn extract_mem_src_id(composite_source_id: &str) -> Option<&str> {
 }
 
 static MEMORY_SYNC_FRONTEND_HANDLE: OnceLock<SubscriptionHandle> = OnceLock::new();
+static MEMORY_SYNC_EMBED_HANDLE: OnceLock<SubscriptionHandle> = OnceLock::new();
 
 /// Register a lightweight bridge that translates lower-level ingestion events
-/// into the coarse sync-stage stream the frontend consumes.
-pub fn register_sync_stage_bridge() {
+/// into the coarse sync-stage stream the frontend consumes, and a post-sync
+/// embed trigger that kicks off batch embedding after sync completion.
+pub fn register_sync_stage_bridge(config: &crate::openhuman::config::Config) {
     if MEMORY_SYNC_FRONTEND_HANDLE.get().is_some() {
         return;
     }
@@ -142,6 +144,45 @@ pub fn register_sync_stage_bridge() {
             );
         }
     }
+
+    // Trigger batch embedding when a sync completes. Extract no longer embeds
+    // inline — the backfill pass picks up all un-embedded chunks in large
+    // batches (up to 1000 items per API call).
+    if MEMORY_SYNC_EMBED_HANDLE.get().is_none() {
+        if let Some(handle) = subscribe_global(Arc::new(SyncCompleteEmbedTrigger {
+            config: config.clone(),
+        })) {
+            let _ = MEMORY_SYNC_EMBED_HANDLE.set(handle);
+            log::debug!("[event_bus] sync-complete embed trigger registered");
+        }
+    }
+}
+
+/// Triggers a `ReembedBackfill` chain when a sync completes so that all
+/// chunks admitted during the sync get their embeddings in one large batch
+/// pass (up to 1000 items per API call, ~1M tokens).
+struct SyncCompleteEmbedTrigger {
+    config: crate::openhuman::config::Config,
+}
+
+#[async_trait]
+impl EventHandler for SyncCompleteEmbedTrigger {
+    fn name(&self) -> &str {
+        "memory::sync_complete_embed_trigger"
+    }
+
+    fn domains(&self) -> Option<&[&str]> {
+        Some(&["memory"])
+    }
+
+    async fn handle(&self, event: &DomainEvent) {
+        if let DomainEvent::MemorySyncStageChanged { stage, .. } = event {
+            if stage == "completed" {
+                log::debug!("[memory-sync] sync completed — triggering batch embedding backfill");
+                crate::openhuman::memory_queue::ensure_reembed_backfill(&self.config);
+            }
+        }
+    }
 }
 
 struct MemorySyncStageBridge;

src/openhuman/memory_queue/handlers/mod.rs

@@ -13,6 +13,7 @@ use anyhow::{Context, Result};
 
 use crate::openhuman::config::Config;
 use crate::openhuman::memory::tree_source::get_or_create_source_tree;
+use crate::openhuman::memory_queue::ensure_reembed_backfill;
 use crate::openhuman::memory_queue::store;
 use crate::openhuman::memory_queue::types::{
     AppendBufferPayload, AppendTarget, ExtractChunkPayload, FlushStalePayload, Job, JobKind,
@@ -35,10 +36,9 @@ const L0_DEFAULT_FLUSH_AGE_SECS: i64 = 60 * 60;
 
 /// Maximum `extract_chunk` jobs to coalesce in one worker tick.
 ///
-/// Scoring/extraction still runs per chunk, but kept chunks share one
-/// `Embedder::embed_batch` call. This is the high-volume memory-sync path:
-/// Slack/Gmail/etc. enqueue many extract jobs, and without coalescing each
-/// kept chunk produces its own embedding request.
+/// Scoring/extraction runs per chunk; embedding is deferred to a post-sync
+/// `ReembedBackfill` pass that maximizes the batch API (up to 1000 items).
+/// Coalescing still reduces per-job transaction overhead.
 pub(crate) const EXTRACT_EMBED_BATCH: usize = 32;
 
 /// Derive the tree scope from a source_id. For GitHub per-item ids like
@@ -184,8 +184,12 @@ async fn handle_extract(config: &Config, job: &Job) -> Result<JobOutcome> {
         .1
 }
 
-/// Handle a claimed run of `extract_chunk` jobs, batching the embedding
-/// sub-step while preserving per-job outcomes for worker settlement.
+/// Handle a claimed run of `extract_chunk` jobs.
+///
+/// Scoring runs per-chunk; embedding is **deferred** to a post-sync
+/// `ReembedBackfill` pass that maximizes the batch API (up to 1000 items
+/// per request, ~1M tokens). After all chunks are finalized, a backfill
+/// is triggered so the embedding pass starts promptly.
 pub async fn handle_extract_batch(
     config: &Config,
     jobs: &[Job],
@@ -201,29 +205,27 @@ pub async fn handle_extract_batch(
         }
     }
 
-    attach_batched_embeddings(config, &mut prepared).await;
-
-    for mut item in prepared {
+    let mut any_admitted = false;
+    for item in prepared {
         let job = item.job.clone();
-        let result = if let Some(err) = item.embedding_error.take() {
-            Err(err)
-        } else {
-            finalize_extract(config, item)
-        };
+        if item.result.kept {
+            any_admitted = true;
+        }
+        let result = finalize_extract(config, item);
         outcomes.push((job, result));
     }
 
+    if any_admitted {
+        ensure_reembed_backfill(config);
+    }
+
     Ok(outcomes)
 }
 
 struct PreparedExtract {
     job: Job,
     chunk: Chunk,
-    body: String,
     result: score::ScoreResult,
-    embedding: Option<Vec<f32>>,
-    embedding_error: Option<anyhow::Error>,
-    embedding_retry_needed: bool,
 }
 
 async fn prepare_extract(config: &Config, job: &Job) -> Result<Option<PreparedExtract>> {
@@ -238,15 +240,13 @@ async fn prepare_extract(config: &Config, job: &Job) -> Result<Option<PreparedEx
     };
 
     // Read the full body from disk (the `content` column in SQLite holds a
-    // ≤500-char preview after the MD-on-disk migration). Both the scorer and
-    // the embedder need the complete text so extraction and semantic indexing
-    // operate over the full chunk body, not a truncated preview.
+    // ≤500-char preview after the MD-on-disk migration). The scorer needs
+    // the complete text so extraction operates over the full chunk body.
     let body = content_read::read_chunk_body(config, &chunk.id)
         .with_context(|| format!("read full body for extract chunk_id={}", chunk.id))?;
-    // Score a clone of the chunk with the full body swapped in.
     let chunk_with_body = {
         let mut c = chunk.clone();
-        c.content = body.clone();
+        c.content = body;
         c
     };
 
@@ -264,129 +264,12 @@ async fn prepare_extract(config: &Config, job: &Job) -> Result<Option<PreparedEx
     Ok(Some(PreparedExtract {
         job: job.clone(),
         chunk,
-        body,
         result,
-        embedding: None,
-        embedding_error: None,
-        embedding_retry_needed: false,
     }))
 }
 
-async fn attach_batched_embeddings(config: &Config, prepared: &mut [PreparedExtract]) {
-    let kept_count = prepared.iter().filter(|item| item.result.kept).count();
-    if kept_count == 0 {
-        return;
-    }
-
-    // #002 (FR-002): when no usable embeddings provider is configured the
-    // write path returns None instead of an InertEmbedder — we SKIP embedding
-    // rather than writing fake all-zero vectors.
-    let embedder = match build_write_embedder(config).context("build embedder in extract handler") {
-        Ok(Some(embedder)) => embedder,
-        Ok(None) => {
-            for item in prepared.iter().filter(|item| item.result.kept) {
-                log::warn!(
-                    "[memory::jobs] extract chunk_id={} — embeddings unavailable, \
-                     skipping embed (semantic recall degraded)",
-                    item.chunk.id
-                );
-            }
-            return;
-        }
-        Err(e) => {
-            for item in prepared.iter_mut().filter(|item| item.result.kept) {
-                let failure = crate::openhuman::memory_tree::health::classify_embed_error(&e);
-                item.embedding_retry_needed = !failure.is_unrecoverable();
-                log::warn!(
-                    "[memory::jobs] extract chunk_id={} — embedder build failed; \
-                     continuing without vector retry_needed={} err={e:#}",
-                    item.chunk.id,
-                    item.embedding_retry_needed
-                );
-            }
-            return;
-        }
-    };
-
-    let kept_indices: Vec<usize> = prepared
-        .iter()
-        .enumerate()
-        .filter_map(|(idx, item)| item.result.kept.then_some(idx))
-        .collect();
-    let texts: Vec<&str> = kept_indices
-        .iter()
-        .map(|idx| prepared[*idx].body.as_str())
-        .collect();
-
-    log::debug!(
-        "[memory::jobs] extract embedding batch: kept_chunks={} provider={}",
-        texts.len(),
-        embedder.name()
-    );
-    let vectors = embedder.embed_batch(&texts).await;
-    if vectors.len() != kept_indices.len() {
-        let err = anyhow::anyhow!(
-            "extract embed_batch returned {} results for {} texts",
-            vectors.len(),
-            kept_indices.len()
-        );
-        let failure = crate::openhuman::memory_tree::health::classify_embed_error(&err);
-        for idx in kept_indices {
-            prepared[idx].embedding_retry_needed = !failure.is_unrecoverable();
-            log::warn!(
-                "[memory::jobs] extract chunk_id={} — batch embedding contract failed; \
-                 continuing without vector retry_needed={} err={err:#}",
-                prepared[idx].chunk.id,
-                prepared[idx].embedding_retry_needed
-            );
-        }
-        return;
-    }
-
-    for (idx, vector_result) in kept_indices.into_iter().zip(vectors) {
-        match vector_result {
-            Ok(vector) => {
-                if let Err(e) = pack_checked(&vector).with_context(|| {
-                    format!(
-                        "validate embedding dims for chunk_id={}",
-                        prepared[idx].chunk.id
-                    )
-                }) {
-                    let failure = crate::openhuman::memory_tree::health::classify_embed_error(&e);
-                    prepared[idx].embedding_retry_needed = !failure.is_unrecoverable();
-                    log::warn!(
-                        "[memory::jobs] extract chunk_id={} — embedding dim validation failed; \
-                         continuing without vector retry_needed={} err={e:#}",
-                        prepared[idx].chunk.id,
-                        prepared[idx].embedding_retry_needed
-                    );
-                } else {
-                    prepared[idx].embedding = Some(vector);
-                    crate::openhuman::memory_tree::health::clear_semantic_recall_degraded();
-                }
-            }
-            Err(e) => {
-                let failure = crate::openhuman::memory_tree::health::classify_embed_error(&e);
-                prepared[idx].embedding_retry_needed = !failure.is_unrecoverable();
-                log::warn!(
-                    "[memory::jobs] extract chunk_id={} — embedding failed; \
-                     continuing without vector retry_needed={} err={e:#}",
-                    prepared[idx].chunk.id,
-                    prepared[idx].embedding_retry_needed
-                );
-            }
-        }
-    }
-}
-
 fn finalize_extract(config: &Config, item: PreparedExtract) -> Result<JobOutcome> {
-    let PreparedExtract {
-        chunk,
-        result,
-        embedding: chunk_embedding,
-        embedding_retry_needed,
-        ..
-    } = item;
+    let PreparedExtract { chunk, result, .. } = item;
     // Build follow-up job payloads before opening the tx — construction is
     // cheap and doesn't require a database connection. The two jobs are
     // enqueued inside the SAME transaction that commits the lifecycle update,
@@ -422,7 +305,6 @@ fn finalize_extract(config: &Config, item: PreparedExtract) -> Result<JobOutcome
         Some(format!("chunk {}", &chunk.id[..chunk.id.len().min(16)])),
     );
 
-    let active_sig = chunk_store::tree_active_signature(config);
     let did_enqueue_source = chunk_store::with_connection(config, |conn| {
         let tx = conn.unchecked_transaction()?;
         score::persist_score_tx(
@@ -439,20 +321,6 @@ fn finalize_extract(config: &Config, item: PreparedExtract) -> Result<JobOutcome
                       WHERE id = ?2",
                 rusqlite::params![chunk_store::CHUNK_STATUS_ADMITTED, chunk.id],
             )?;
-            // #1574 write-side cutover: persist the embedding to the
-            // per-model `mem_tree_chunk_embeddings` sidecar at the active
-            // signature, inside THIS tx so it commits atomically with the
-            // lifecycle / score / job-enqueue writes. The legacy
-            // `mem_tree_chunks.embedding` column is no longer written
-            // (left intact for the §7 one-shot migration to read).
-            if let Some(emb) = chunk_embedding.as_deref() {
-                chunk_store::set_chunk_embedding_for_signature_tx(
-                    &tx,
-                    &chunk.id,
-                    &active_sig,
-                    emb,
-                )?;
-            }
         } else {
             tx.execute(
                 "UPDATE mem_tree_chunks
@@ -516,14 +384,9 @@ fn finalize_extract(config: &Config, item: PreparedExtract) -> Result<JobOutcome
         }
     }
 
-    // Signal workers after the tx commits (no atomicity requirement on signaling).
     if did_enqueue_source {
         super::worker::wake_workers();
     }
-    if result.kept && embedding_retry_needed {
-        crate::openhuman::memory_queue::ensure_reembed_backfill(config);
-        super::worker::wake_workers();
-    }
 
     Ok(JobOutcome::Done)
 }
@@ -787,11 +650,11 @@ async fn handle_flush_stale(config: &Config, job: &Job) -> Result<JobOutcome> {
     Ok(JobOutcome::Done)
 }
 
-/// Texts per `ReembedBackfill` run. Bounded so one run holds the global
-/// single-LLM-slot (the job is `is_llm_bound`) for a predictable spell —
-/// the laptop-RAM safety the local-LLM-load rule requires. The chain
-/// self-continues via `Defer` until no rows remain.
-const REEMBED_BACKFILL_BATCH: usize = 16;
+/// Texts per `ReembedBackfill` run. This is now the **primary** embedding
+/// path (extract no longer embeds inline). Sized to maximize the batch API
+/// (Voyage: 1000 items, ~1M tokens per request). The `embed_batch_via_provider`
+/// layer handles sub-batching into API-safe chunks internally.
+const REEMBED_BACKFILL_BATCH: usize = 1000;
 /// Delay before the deferred chain revisits this same job row.
 const REEMBED_BACKFILL_REVISIT_MS: i64 = 750;
 

src/openhuman/memory_tree/score/embed/mod.rs

@@ -104,16 +104,49 @@ pub(crate) fn check_embed_dim(v: Vec<f32>, label: &str) -> Result<Vec<f32>> {
     Ok(v)
 }
 
-/// Batch-embed `texts` through a unified [`EmbeddingProvider`] in a single
-/// request and adapt the result to the [`Embedder::embed_batch`] contract:
-/// one dimension-checked [`Result`] per input position.
+/// Voyage batch API limits (conservative estimates).
+const MAX_BATCH_ITEMS: usize = 1000;
+const MAX_BATCH_TOKENS: usize = 1_000_000;
+const CHARS_PER_TOKEN_ESTIMATE: usize = 4;
+
+fn estimate_tokens(text: &str) -> usize {
+    (text.len() + CHARS_PER_TOKEN_ESTIMATE - 1) / CHARS_PER_TOKEN_ESTIMATE
+}
+
+/// Split `texts` into sub-batches that respect the batch API limits:
+/// at most `MAX_BATCH_ITEMS` items per batch and at most
+/// `MAX_BATCH_TOKENS` estimated tokens per batch.
+fn split_into_sub_batches<'a>(texts: &[&'a str]) -> Vec<Vec<&'a str>> {
+    let mut batches: Vec<Vec<&'a str>> = Vec::new();
+    let mut current: Vec<&'a str> = Vec::new();
+    let mut current_tokens: usize = 0;
+
+    for &text in texts {
+        let tokens = estimate_tokens(text);
+        if !current.is_empty()
+            && (current.len() >= MAX_BATCH_ITEMS || current_tokens + tokens > MAX_BATCH_TOKENS)
+        {
+            batches.push(std::mem::take(&mut current));
+            current_tokens = 0;
+        }
+        current.push(text);
+        current_tokens += tokens;
+    }
+    if !current.is_empty() {
+        batches.push(current);
+    }
+    batches
+}
+
+/// Batch-embed `texts` through a unified [`EmbeddingProvider`], splitting
+/// into sub-batches that respect the batch API limits (1000 items, ~1M
+/// tokens per request).
 ///
-/// On a wholesale batch failure (network / auth) **or** a length-contract
-/// violation from the provider, this falls back to per-text
-/// [`EmbeddingProvider::embed_one`] so a single transient blip cannot fail —
-/// and, in the backfill, *tombstone* — every row in the batch. That preserves
-/// the resilience of the original per-item loop while keeping the happy-path
-/// single-round-trip win.
+/// Each sub-batch is sent as a single provider `embed()` call. On a
+/// wholesale batch failure **or** a length-contract violation, the failing
+/// sub-batch falls back to per-text [`EmbeddingProvider::embed_one`] so a
+/// single transient blip cannot fail — and, in the backfill, *tombstone* —
+/// every row in the batch.
 pub(crate) async fn embed_batch_via_provider(
     inner: &dyn crate::openhuman::embeddings::EmbeddingProvider,
     label: &str,
@@ -122,14 +155,37 @@ pub(crate) async fn embed_batch_via_provider(
     if texts.is_empty() {
         return Vec::new();
     }
+
+    let sub_batches = split_into_sub_batches(texts);
     log::debug!(
-        "[memory_tree::embed::{label}] embed_batch:enter texts={}",
-        texts.len()
+        "[memory_tree::embed::{label}] embed_batch:enter texts={} sub_batches={}",
+        texts.len(),
+        sub_batches.len()
     );
+
+    let mut all_results: Vec<Result<Vec<f32>>> = Vec::with_capacity(texts.len());
+
+    for (batch_idx, batch) in sub_batches.iter().enumerate() {
+        let batch_results = embed_one_sub_batch(inner, label, batch, batch_idx).await;
+        all_results.extend(batch_results);
+    }
+
+    all_results
+}
+
+/// Embed a single sub-batch via the provider, with per-text fallback on
+/// batch failure.
+async fn embed_one_sub_batch(
+    inner: &dyn crate::openhuman::embeddings::EmbeddingProvider,
+    label: &str,
+    texts: &[&str],
+    batch_idx: usize,
+) -> Vec<Result<Vec<f32>>> {
     match inner.embed(texts).await {
         Ok(vectors) if vectors.len() == texts.len() => {
             log::debug!(
-                "[memory_tree::embed::{label}] embed_batch:success collapsed {} texts into one provider call",
+                "[memory_tree::embed::{label}] embed_batch:success sub_batch={batch_idx} \
+                 collapsed {} texts into one provider call",
                 texts.len()
             );
             vectors
@@ -139,8 +195,8 @@ pub(crate) async fn embed_batch_via_provider(
         }
         Ok(vectors) => {
             log::warn!(
-                "[memory_tree::embed::{label}] embed_batch:fallback batch returned {} vectors for {} texts; \
-                 falling back to per-text embedding",
+                "[memory_tree::embed::{label}] embed_batch:fallback sub_batch={batch_idx} \
+                 returned {} vectors for {} texts; falling back to per-text embedding",
                 vectors.len(),
                 texts.len()
             );
@@ -148,8 +204,8 @@ pub(crate) async fn embed_batch_via_provider(
         }
         Err(e) => {
             log::warn!(
-                "[memory_tree::embed::{label}] embed_batch:fallback batch embed failed ({e:#}); \
-                 falling back to per-text embedding"
+                "[memory_tree::embed::{label}] embed_batch:fallback sub_batch={batch_idx} \
+                 batch embed failed ({e:#}); falling back to per-text embedding"
             );
             embed_each_via_provider(inner, label, texts).await
         }