engram_core/
consolidate_handler.rs

1use std::sync::Arc;
2
3use serde::Deserialize;
4use serde_json::{Value, json};
5
6use crate::error::CoreError;
7use crate::server::ServerState;
8
9const MAX_STALE_DAYS: u32 = 3650;
10
11#[derive(Deserialize)]
12struct PreviewParams {
13    stale_days: Option<u32>,
14    min_score: Option<f64>,
15}
16
17pub async fn handle_preview(state: &Arc<ServerState>, params: Value) -> Result<Value, CoreError> {
18    let parsed: PreviewParams = serde_json::from_value(params)
19        .map_err(|error| CoreError::DispatchError(error.to_string()))?;
20    let stale_days = parsed
21        .stale_days
22        .unwrap_or(state.config.consolidation.stale_days);
23    let min_score = parsed
24        .min_score
25        .unwrap_or(state.config.consolidation.min_score);
26    validate_consolidation_params(stale_days, min_score)?;
27    let state_clone = Arc::clone(state);
28    let result = tokio::task::spawn_blocking(move || {
29        let database = state_clone.database.lock().unwrap();
30        let preview = engram_consolidate::preview(&database, stale_days, min_score)?;
31        Ok::<_, CoreError>(json!({
32            "duplicates": preview.duplicates.len(),
33            "stale": preview.stale.len(),
34            "garbage": preview.garbage.len(),
35            "duplicate_groups": serialize_duplicate_groups(&preview.duplicates),
36            "stale_ids": preview.stale,
37            "garbage_ids": preview.garbage,
38        }))
39    })
40    .await
41    .map_err(|error| CoreError::SocketError(error.to_string()))??;
42    Ok(result)
43}
44
45pub async fn handle_analyze(state: &Arc<ServerState>, params: Value) -> Result<Value, CoreError> {
46    let parsed: PreviewParams = serde_json::from_value(params)
47        .map_err(|error| CoreError::DispatchError(error.to_string()))?;
48    let stale_days = parsed
49        .stale_days
50        .unwrap_or(state.config.consolidation.stale_days);
51    let min_score = parsed
52        .min_score
53        .unwrap_or(state.config.consolidation.min_score);
54    validate_consolidation_params(stale_days, min_score)?;
55    let state_clone = Arc::clone(state);
56    let result = tokio::task::spawn_blocking(move || {
57        let database = state_clone.database.lock().unwrap();
58        let preview = engram_consolidate::preview(&database, stale_days, min_score)?;
59        let text_gen_ref = state_clone
60            .text_generator
61            .as_deref()
62            .map(|generator| generator as &dyn engram_llm_client::TextGenerator);
63        let analysis = engram_consolidate::analyze(&database, &preview, text_gen_ref)?;
64        Ok::<_, CoreError>(json!({
65            "analyzed_count": analysis.analyzed_count,
66            "recommendations": serialize_recommendations(&analysis.recommendations),
67        }))
68    })
69    .await
70    .map_err(|error| CoreError::SocketError(error.to_string()))??;
71    Ok(result)
72}
73
74#[derive(Deserialize)]
75struct ApplyParams {
76    stale_days: Option<u32>,
77    min_score: Option<f64>,
78}
79
80pub async fn handle_apply(state: &Arc<ServerState>, params: Value) -> Result<Value, CoreError> {
81    let parsed: ApplyParams = serde_json::from_value(params)
82        .map_err(|error| CoreError::DispatchError(error.to_string()))?;
83    let stale_days = parsed
84        .stale_days
85        .unwrap_or(state.config.consolidation.stale_days);
86    let min_score = parsed
87        .min_score
88        .unwrap_or(state.config.consolidation.min_score);
89    validate_consolidation_params(stale_days, min_score)?;
90    let state_clone = Arc::clone(state);
91    let result = tokio::task::spawn_blocking(move || {
92        let database = state_clone.database.lock().unwrap();
93        let preview = engram_consolidate::preview(&database, stale_days, min_score)?;
94        let text_gen_ref = state_clone
95            .text_generator
96            .as_deref()
97            .map(|generator| generator as &dyn engram_llm_client::TextGenerator);
98        let analysis = engram_consolidate::analyze(&database, &preview, text_gen_ref)?;
99        let apply_result =
100            engram_consolidate::apply(&database, &analysis.recommendations, "server")?;
101        Ok::<_, CoreError>(json!({
102            "merged": apply_result.merged,
103            "deleted": apply_result.deleted,
104            "archived": apply_result.archived,
105            "kept": apply_result.kept,
106            "errors": apply_result.errors,
107        }))
108    })
109    .await
110    .map_err(|error| CoreError::SocketError(error.to_string()))??;
111    Ok(result)
112}
113
114fn serialize_duplicate_groups(groups: &[engram_consolidate::DuplicateGroup]) -> Vec<Value> {
115    groups
116        .iter()
117        .map(|group| {
118            json!({
119                "primary_id": group.primary_id,
120                "duplicate_ids": group.duplicate_ids,
121                "similarity": group.similarity,
122            })
123        })
124        .collect()
125}
126
127fn serialize_recommendations(recommendations: &[engram_consolidate::Recommendation]) -> Vec<Value> {
128    recommendations
129        .iter()
130        .map(|recommendation| {
131            json!({
132                "action": format_action(&recommendation.action),
133                "confidence": recommendation.confidence,
134                "reasoning": recommendation.reasoning,
135            })
136        })
137        .collect()
138}
139
140fn format_action(action: &engram_consolidate::RecommendedAction) -> Value {
141    match action {
142        engram_consolidate::RecommendedAction::Merge {
143            source_id,
144            target_id,
145        } => json!({ "type": "merge", "source_id": source_id, "target_id": target_id }),
146        engram_consolidate::RecommendedAction::Delete { memory_id } => {
147            json!({ "type": "delete", "memory_id": memory_id })
148        }
149        engram_consolidate::RecommendedAction::Archive { memory_id } => {
150            json!({ "type": "archive", "memory_id": memory_id })
151        }
152        engram_consolidate::RecommendedAction::Keep { memory_id } => {
153            json!({ "type": "keep", "memory_id": memory_id })
154        }
155    }
156}
157
158fn validate_consolidation_params(stale_days: u32, min_score: f64) -> Result<(), CoreError> {
159    if stale_days > MAX_STALE_DAYS {
160        return Err(CoreError::DispatchError(format!(
161            "stale_days exceeds maximum of {MAX_STALE_DAYS}"
162        )));
163    }
164    if !(0.0..=1.0).contains(&min_score) {
165        return Err(CoreError::DispatchError(
166            "min_score must be between 0.0 and 1.0".into(),
167        ));
168    }
169    Ok(())
170}
engram_core/consolidate_handler.rs

engram_core/
consolidate_handler.rs
