engram_core/
consolidate_handler.rs

1use std::sync::Arc;
2
3use serde::Deserialize;
4use serde_json::{Value, json};
5
6use engram_storage::Database;
7
8use crate::error::CoreError;
9use crate::lock_helpers;
10use crate::persistence::hash_string_to_u64;
11use crate::server::ServerState;
12
13const MAX_STALE_DAYS: u32 = 3650;
14const MIN_CONFIDENCE_DEFAULT: f32 = 0.0;
15
16#[derive(Deserialize)]
17struct PreviewParams {
18    stale_days: Option<u32>,
19    min_score: Option<f64>,
20}
21
22pub async fn handle_preview(state: &Arc<ServerState>, params: Value) -> Result<Value, CoreError> {
23    let parsed: PreviewParams = serde_json::from_value(params)
24        .map_err(|error| CoreError::DispatchError(error.to_string()))?;
25    let stale_days = parsed
26        .stale_days
27        .unwrap_or(state.config.consolidation.stale_days);
28    let min_score = parsed
29        .min_score
30        .unwrap_or(state.config.consolidation.min_score);
31    validate_consolidation_params(stale_days, min_score)?;
32    let state_clone = Arc::clone(state);
33    let result = tokio::task::spawn_blocking(move || {
34        let database = lock_helpers::lock_db(&state_clone);
35        let preview = run_preview(&state_clone, &database, stale_days, min_score)?;
36        Ok::<_, CoreError>(json!({
37            "duplicates": preview.duplicates.len(),
38            "stale": preview.stale.len(),
39            "garbage": preview.garbage.len(),
40            "duplicate_groups": serialize_duplicate_groups(&preview.duplicates),
41            "stale_ids": preview.stale,
42            "garbage_ids": preview.garbage,
43        }))
44    })
45    .await
46    .map_err(|error| CoreError::SocketError(error.to_string()))??;
47    Ok(result)
48}
49
50pub async fn handle_analyze(state: &Arc<ServerState>, params: Value) -> Result<Value, CoreError> {
51    let parsed: PreviewParams = serde_json::from_value(params)
52        .map_err(|error| CoreError::DispatchError(error.to_string()))?;
53    let stale_days = parsed
54        .stale_days
55        .unwrap_or(state.config.consolidation.stale_days);
56    let min_score = parsed
57        .min_score
58        .unwrap_or(state.config.consolidation.min_score);
59    validate_consolidation_params(stale_days, min_score)?;
60    let state_clone = Arc::clone(state);
61    let result = tokio::task::spawn_blocking(move || {
62        let database = lock_helpers::lock_db(&state_clone);
63        let preview = run_preview(&state_clone, &database, stale_days, min_score)?;
64        let text_gen_ref = state_clone
65            .text_generator
66            .as_deref()
67            .map(|generator| generator as &dyn engram_llm_client::TextGenerator);
68        let analysis = engram_consolidate::analyze(&database, &preview, text_gen_ref)?;
69        let recommendations = serialize_recommendations(&analysis.recommendations);
70        Ok::<_, CoreError>(json!({
71            "analyzed_count": analysis.analyzed_count,
72            "recommendations": recommendations,
73            "errors": analysis.errors,
74        }))
75    })
76    .await
77    .map_err(|error| CoreError::SocketError(error.to_string()))??;
78    Ok(result)
79}
80
81#[derive(Deserialize)]
82struct ApplyParams {
83    stale_days: Option<u32>,
84    min_score: Option<f64>,
85    min_confidence: Option<f32>,
86}
87
88pub async fn handle_apply(state: &Arc<ServerState>, params: Value) -> Result<Value, CoreError> {
89    let parsed: ApplyParams = serde_json::from_value(params)
90        .map_err(|error| CoreError::DispatchError(error.to_string()))?;
91    let stale_days = parsed
92        .stale_days
93        .unwrap_or(state.config.consolidation.stale_days);
94    let min_score = parsed
95        .min_score
96        .unwrap_or(state.config.consolidation.min_score);
97    let min_confidence = parsed.min_confidence.unwrap_or(MIN_CONFIDENCE_DEFAULT);
98    validate_consolidation_params(stale_days, min_score)?;
99    validate_min_confidence(min_confidence)?;
100    let state_clone = Arc::clone(state);
101    let result = tokio::task::spawn_blocking(move || {
102        let (apply_result, mut errors) = {
103            let database = lock_helpers::lock_db(&state_clone);
104            let preview = run_preview(&state_clone, &database, stale_days, min_score)?;
105            let text_gen_ref = state_clone
106                .text_generator
107                .as_deref()
108                .map(|generator| generator as &dyn engram_llm_client::TextGenerator);
109            let analysis = engram_consolidate::analyze(&database, &preview, text_gen_ref)?;
110            let apply_result = engram_consolidate::apply(
111                &database,
112                &analysis.recommendations,
113                "server",
114                min_confidence,
115            )?;
116            (apply_result, analysis.errors)
117        };
118        let mut indexes = lock_helpers::write_indexes(&state_clone);
119        for id in &apply_result.pruned_ids {
120            let hashed = hash_string_to_u64(id);
121            if indexes.contains(hashed) {
122                indexes.delete(hashed).map_err(CoreError::Hnsw)?;
123            }
124        }
125        errors.extend(apply_result.errors.iter().cloned());
126        Ok::<_, CoreError>(json!({
127            "merged": apply_result.merged,
128            "deleted": apply_result.deleted,
129            "archived": apply_result.archived,
130            "kept": apply_result.kept,
131            "errors": errors,
132        }))
133    })
134    .await
135    .map_err(|error| CoreError::SocketError(error.to_string()))??;
136    Ok(result)
137}
138
139fn run_preview(
140    state: &ServerState,
141    database: &Database,
142    stale_days: u32,
143    min_score: f64,
144) -> Result<engram_consolidate::PreviewResult, CoreError> {
145    let preview = engram_consolidate::preview(
146        database,
147        stale_days,
148        min_score,
149        state.config.consolidation.fts_similarity_floor,
150    )?;
151    Ok(preview)
152}
153
154fn serialize_duplicate_groups(groups: &[engram_consolidate::DuplicateGroup]) -> Vec<Value> {
155    groups
156        .iter()
157        .map(|group| {
158            json!({
159                "primary_id": group.primary_id,
160                "duplicate_ids": group.duplicate_ids,
161                "similarity": group.similarity,
162                "match_type": group.match_type.as_str(),
163            })
164        })
165        .collect()
166}
167
168fn serialize_recommendations(recommendations: &[engram_consolidate::Recommendation]) -> Vec<Value> {
169    recommendations
170        .iter()
171        .map(|recommendation| {
172            json!({
173                "action": format_action(&recommendation.action),
174                "confidence": recommendation.confidence,
175                "reasoning": recommendation.reasoning,
176            })
177        })
178        .collect()
179}
180
181fn format_action(action: &engram_consolidate::RecommendedAction) -> Value {
182    match action {
183        engram_consolidate::RecommendedAction::Merge {
184            source_id,
185            target_id,
186        } => json!({ "type": "merge", "source_id": source_id, "target_id": target_id }),
187        engram_consolidate::RecommendedAction::Delete { memory_id } => {
188            json!({ "type": "delete", "memory_id": memory_id })
189        }
190        engram_consolidate::RecommendedAction::Archive { memory_id } => {
191            json!({ "type": "archive", "memory_id": memory_id })
192        }
193        engram_consolidate::RecommendedAction::Keep { memory_id } => {
194            json!({ "type": "keep", "memory_id": memory_id })
195        }
196    }
197}
198
199fn validate_consolidation_params(stale_days: u32, min_score: f64) -> Result<(), CoreError> {
200    if stale_days > MAX_STALE_DAYS {
201        return Err(CoreError::DispatchError(format!(
202            "stale_days exceeds maximum of {MAX_STALE_DAYS}"
203        )));
204    }
205    if !(0.0..=1.0).contains(&min_score) {
206        return Err(CoreError::DispatchError(
207            "min_score must be between 0.0 and 1.0".into(),
208        ));
209    }
210    Ok(())
211}
212
213fn validate_min_confidence(min_confidence: f32) -> Result<(), CoreError> {
214    if !(0.0..=1.0).contains(&min_confidence) {
215        return Err(CoreError::DispatchError(
216            "min_confidence must be between 0.0 and 1.0".into(),
217        ));
218    }
219    Ok(())
220}
221
222#[cfg(test)]
223mod tests {
224    use super::*;
225
226    #[test]
227    fn min_confidence_in_range_is_accepted() {
228        assert!(validate_min_confidence(0.0).is_ok());
229        assert!(validate_min_confidence(0.5).is_ok());
230        assert!(validate_min_confidence(1.0).is_ok());
231    }
232
233    #[test]
234    fn min_confidence_out_of_range_is_rejected() {
235        assert!(matches!(
236            validate_min_confidence(-0.1),
237            Err(CoreError::DispatchError(_))
238        ));
239        assert!(matches!(
240            validate_min_confidence(1.1),
241            Err(CoreError::DispatchError(_))
242        ));
243    }
244
245    #[test]
246    fn min_confidence_nan_is_rejected() {
247        assert!(matches!(
248            validate_min_confidence(f32::NAN),
249            Err(CoreError::DispatchError(_))
250        ));
251    }
252}
engram_core/consolidate_handler.rs

engram_core/
consolidate_handler.rs
