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

engram_core/
reembed_handler.rs
