engram_core/
reembed_handler.rs

1use std::sync::Arc;
2
3use serde::Deserialize;
4use serde_json::{Value, json};
5
6use engram_embeddings::ThreeFieldEmbedding;
7use engram_storage::Memory;
8
9use crate::error::CoreError;
10use crate::lock_helpers;
11use crate::persistence::{deterministic_rng, f32_vec_to_bytes, hash_string_to_u64};
12use crate::server::ServerState;
13
14#[derive(Deserialize, Default)]
15struct ReembedParams {
16    /// Reserved for future safety thresholds (e.g., refuse if memory count
17    /// exceeds a configured limit). Currently a no-op placeholder per
18    /// ADR 2026-05-14-voyage-4-migration-via-reembed-cli §Decision step 5.
19    #[serde(default)]
20    force: bool,
21}
22
23pub async fn handle(state: &Arc<ServerState>, params: Value) -> Result<Value, CoreError> {
24    let parsed: ReembedParams = serde_json::from_value(params).unwrap_or_default();
25    let _ = parsed.force;
26
27    state.embedder.clear_cache();
28
29    let state_clone = Arc::clone(state);
30    tokio::task::spawn_blocking(move || run_reembed(state_clone))
31        .await
32        .map_err(|error| CoreError::SocketError(error.to_string()))?
33}
34
35fn run_reembed(state: Arc<ServerState>) -> Result<Value, CoreError> {
36    let memories = {
37        let database = lock_helpers::lock_db(&state);
38        database.list_all_memories().map_err(CoreError::Storage)?
39    };
40    let total = memories.len();
41    let model = state.embedding_provider.model_name().to_string();
42
43    let mut succeeded = 0usize;
44    let mut failed = 0usize;
45    let mut first_failure_reported = false;
46
47    for memory in &memories {
48        match reembed_one(&state, memory) {
49            Ok(()) => succeeded += 1,
50            Err(ReembedError::Recoverable(error)) => {
51                if !first_failure_reported {
52                    eprintln!(
53                        "reembed: first recoverable failure on memory {}: {error}. \
54                         Subsequent failures suppressed; rerun reembed after fixing \
55                         the cause to retry only memories left with indexed=0.",
56                        memory.id
57                    );
58                    first_failure_reported = true;
59                }
60                mark_for_retry(&state, &memory.id)?;
61                failed += 1;
62            }
63            Err(ReembedError::Fatal(error)) => return Err(error),
64        }
65    }
66
67    if failed == 0 {
68        let database = lock_helpers::lock_db(&state);
69        crate::migrations::embedding_model_v1::record(&database, &model)?;
70    }
71
72    Ok(json!({
73        "total": total,
74        "succeeded": succeeded,
75        "failed": failed,
76        "model": model,
77    }))
78}
79
80enum ReembedError {
81    /// Provider or HNSW transient failure — memory stays in the DB with
82    /// stale embeddings; `indexed=0` flags it for retry by background
83    /// reindex (server mode) or a subsequent reembed invocation. Carries
84    /// the underlying error so the orchestrator can log the first failure
85    /// (Voyage rate limit, auth, etc.) for the user.
86    Recoverable(CoreError),
87    /// Database write failure or invariant violation — propagate.
88    Fatal(CoreError),
89}
90
91fn reembed_one(state: &Arc<ServerState>, memory: &Memory) -> Result<(), ReembedError> {
92    let provider = state.embedding_provider.as_ref();
93    let text_gen = state
94        .text_generator
95        .as_deref()
96        .map(|generator| generator as &dyn engram_llm_client::TextGenerator);
97
98    let embedding = state
99        .embedder
100        .embed_fields(
101            &memory.context,
102            &memory.action,
103            &memory.result,
104            provider,
105            text_gen,
106        )
107        .map_err(
108            |engram_embeddings::EmbeddingError::ProviderError(api_error)| {
109                ReembedError::Recoverable(CoreError::Api(api_error))
110            },
111        )?;
112
113    replace_in_hnsw(state, &memory.id, &embedding).map_err(ReembedError::Recoverable)?;
114    persist_embeddings(state, &memory.id, &embedding).map_err(ReembedError::Fatal)?;
115
116    Ok(())
117}
118
119fn replace_in_hnsw(
120    state: &Arc<ServerState>,
121    memory_id: &str,
122    embedding: &ThreeFieldEmbedding,
123) -> Result<(), CoreError> {
124    let mut indexes = lock_helpers::write_indexes(state);
125    let hashed = hash_string_to_u64(memory_id);
126    if indexes.contains(hashed) {
127        indexes.delete(hashed).map_err(CoreError::Hnsw)?;
128    }
129    let rng_value = deterministic_rng(hashed);
130    indexes.insert_atomic(hashed, memory_id, embedding, rng_value)
131}
132
133fn persist_embeddings(
134    state: &Arc<ServerState>,
135    memory_id: &str,
136    embedding: &ThreeFieldEmbedding,
137) -> Result<(), CoreError> {
138    let database = lock_helpers::lock_db(state);
139    database
140        .set_memory_embeddings(
141            memory_id,
142            &f32_vec_to_bytes(&embedding.context),
143            &f32_vec_to_bytes(&embedding.action),
144            &f32_vec_to_bytes(&embedding.result),
145        )
146        .map_err(CoreError::Storage)?;
147    database
148        .set_memory_indexed(memory_id, true)
149        .map_err(CoreError::Storage)
150}
151
152fn mark_for_retry(state: &Arc<ServerState>, memory_id: &str) -> Result<(), CoreError> {
153    let database = lock_helpers::lock_db(state);
154    database
155        .set_memory_indexed(memory_id, false)
156        .map_err(CoreError::Storage)
157}
engram_core/reembed_handler.rs

engram_core/
reembed_handler.rs
