From 08cbd9813d5e44c59d5d40024121a9e1d5aa5673 Mon Sep 17 00:00:00 2001 From: Steven Enamakel <31011319+senamakel@users.noreply.github.com> Date: Sat, 6 Jun 2026 07:17:30 -0400 Subject: [PATCH] feat(memory): defer embeddings to post-sync batch pass (#3428) --- src/core/jsonrpc.rs | 2 +- src/openhuman/channels/runtime/startup.rs | 2 +- src/openhuman/memory/sync.rs | 45 ++++- src/openhuman/memory_queue/handlers/mod.rs | 195 +++---------------- src/openhuman/memory_tree/score/embed/mod.rs | 88 +++++++-- 5 files changed, 146 insertions(+), 186 deletions(-) diff --git a/src/core/jsonrpc.rs b/src/core/jsonrpc.rs index 0021fba35..e8fe50e28 100644 --- a/src/core/jsonrpc.rs +++ b/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(), ) { diff --git a/src/openhuman/channels/runtime/startup.rs b/src/openhuman/channels/runtime/startup.rs index b446c56c1..22a35e19e 100644 --- a/src/openhuman/channels/runtime/startup.rs +++ b/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). diff --git a/src/openhuman/memory/sync.rs b/src/openhuman/memory/sync.rs index e1eaa28d2..745108cfd 100644 --- a/src/openhuman/memory/sync.rs +++ b/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 = OnceLock::new(); +static MEMORY_SYNC_EMBED_HANDLE: OnceLock = 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; diff --git a/src/openhuman/memory_queue/handlers/mod.rs b/src/openhuman/memory_queue/handlers/mod.rs index a091d0505..2c6641aa1 100644 --- a/src/openhuman/memory_queue/handlers/mod.rs +++ b/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 { .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>, - embedding_error: Option, - embedding_retry_needed: bool, } async fn prepare_extract(config: &Config, job: &Job) -> Result> { @@ -238,15 +240,13 @@ async fn prepare_extract(config: &Config, job: &Job) -> Result Result 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 = 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 { - 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 Result Result Result { 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; diff --git a/src/openhuman/memory_tree/score/embed/mod.rs b/src/openhuman/memory_tree/score/embed/mod.rs index 00a783085..aaf5df0fe 100644 --- a/src/openhuman/memory_tree/score/embed/mod.rs +++ b/src/openhuman/memory_tree/score/embed/mod.rs @@ -104,16 +104,49 @@ pub(crate) fn check_embed_dim(v: Vec, label: &str) -> Result> { 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> { + let mut batches: Vec> = 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>> = 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>> { 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 }