engine/

compress.rs

//! CE-10b: Memory compression via DBSCAN-based semantic clustering.
//!
//! Clusters old, low-importance memories using DBSCAN (eps = 0.88, more
//! permissive than the consolidation default of 0.92) and synthesises a single
//! summary memory for each cluster by concatenating member content.  Originals
//! are soft-deprecated (expires_at = now + 30 days) rather than hard-deleted.
//!
//! This is the *manual* compression API — see `POST /v1/agents/{id}/compress`.
//! It is complementary to the background DBSCAN consolidation (CE-6).

use std::collections::{HashMap, HashSet};
use std::time::{SystemTime, UNIX_EPOCH};

use common::{Memory, MemoryType};

use crate::consolidate::cosine_sim;

/// Memories above this importance threshold are not candidates for compression.
/// High-importance memories should be preserved intact.
const COMPRESS_MAX_IMPORTANCE: f32 = 0.6;

/// DBSCAN epsilon for compression clustering (more permissive than CE-6's 0.92).
const COMPRESS_EPSILON: f32 = 0.88;

/// Minimum cluster size; singletons are never compressed.
const COMPRESS_MIN_SAMPLES: usize = 2;

/// Soft-deprecation window for compressed originals (30 days).
const SOFT_DEPRECATION_SECS: u64 = 30 * 86_400;

/// A `(Memory, embedding)` pair — used as input/output of compression.
pub type MemoryEmbedPair = (Memory, Vec<f32>);

/// Result returned by [`compress_memories`].
#[derive(Debug, Default)]
pub struct CompressResult {
    /// Total memories evaluated as compression candidates.
    pub memories_scanned: usize,
    /// Number of DBSCAN clusters formed.
    pub clusters_found: usize,
    /// Summary memories created (one per cluster).
    pub summaries_created: usize,
    /// Original cluster members that were soft-deprecated.
    pub originals_deprecated: usize,
    /// IDs of the newly created summary memories.
    pub summary_ids: Vec<String>,
    /// IDs of the soft-deprecated originals.
    pub deprecated_ids: Vec<String>,
}

/// Compress an agent's memories by clustering semantically similar
/// low-importance memories into summary memories.
///
/// # Returns
/// `(new_summary_memories, updated_original_memories, result)`
///
/// The caller is responsible for persisting all returned memories to storage.
/// Summary memories use centroid embeddings of the cluster as their vector.
pub fn compress_memories(
    memories: &[MemoryEmbedPair],
) -> (Vec<MemoryEmbedPair>, Vec<MemoryEmbedPair>, CompressResult) {
    let now_secs = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs();

    // Only consider unexpired, low-importance memories as candidates.
    let candidate_indices: Vec<usize> = memories
        .iter()
        .enumerate()
        .filter(|(_, (m, _))| m.importance <= COMPRESS_MAX_IMPORTANCE && m.expires_at.is_none())
        .map(|(i, _)| i)
        .collect();

    let mut result = CompressResult {
        memories_scanned: candidate_indices.len(),
        ..Default::default()
    };

    if candidate_indices.len() < COMPRESS_MIN_SAMPLES {
        return (Vec::new(), Vec::new(), result);
    }

    // --- DBSCAN (eps = 0.88) on candidate subset -------------------------

    let n = candidate_indices.len();

    // Build neighbour lists (symmetric)
    let mut neighbors: HashMap<usize, Vec<usize>> = HashMap::new();
    for p in 0..n {
        for q in (p + 1)..n {
            let i = candidate_indices[p];
            let j = candidate_indices[q];
            let sim = cosine_sim(&memories[i].1, &memories[j].1);
            if sim >= COMPRESS_EPSILON {
                neighbors.entry(p).or_default().push(q);
                neighbors.entry(q).or_default().push(p);
            }
        }
    }

    // Identify core points (≥ min_samples - 1 neighbours).
    let min_nb = COMPRESS_MIN_SAMPLES.saturating_sub(1).max(1);
    let core: HashSet<usize> = (0..n)
        .filter(|p| neighbors.get(p).map_or(0, |v| v.len()) >= min_nb)
        .collect();

    // Expand clusters.
    let mut visited: HashSet<usize> = HashSet::new();
    let mut clusters: Vec<Vec<usize>> = Vec::new();

    for &cp in &core {
        if visited.contains(&cp) {
            continue;
        }
        let mut cluster = Vec::new();
        let mut stack = vec![cp];
        while let Some(node) = stack.pop() {
            if visited.insert(node) {
                cluster.push(node);
                if let Some(nbrs) = neighbors.get(&node) {
                    for &nb in nbrs {
                        if core.contains(&nb) && !visited.contains(&nb) {
                            stack.push(nb);
                        }
                    }
                }
            }
        }
        if cluster.len() >= COMPRESS_MIN_SAMPLES {
            clusters.push(cluster);
        }
    }

    result.clusters_found = clusters.len();
    if clusters.is_empty() {
        return (Vec::new(), Vec::new(), result);
    }

    // --- Build summary memories and deprecate originals -------------------

    let expires_at = now_secs + SOFT_DEPRECATION_SECS;
    let mut summaries: Vec<(Memory, Vec<f32>)> = Vec::new();
    let mut deprecated: Vec<(Memory, Vec<f32>)> = Vec::new();

    for cluster in &clusters {
        // Collect member memories.
        let members: Vec<&(Memory, Vec<f32>)> = cluster
            .iter()
            .map(|&p| &memories[candidate_indices[p]])
            .collect();

        // Concatenate content into a summary.
        let summary_content = members
            .iter()
            .map(|(m, _)| m.content.as_str())
            .collect::<Vec<_>>()
            .join(" | ");

        // Use max importance and oldest created_at from the cluster.
        let max_importance = members
            .iter()
            .map(|(m, _)| m.importance)
            .fold(f32::NEG_INFINITY, f32::max);
        let oldest_created_at = members
            .iter()
            .map(|(m, _)| m.created_at)
            .min()
            .unwrap_or(now_secs);

        let agent_id = members[0].0.agent_id.clone();
        let summary_id = format!(
            "mem_compress_{:x}",
            now_secs ^ (cluster[0] as u64 * 0x9e3779b97f4a7c15)
        );

        // Compute centroid embedding for the summary (average of cluster embeddings).
        let dim = members[0].1.len();
        let centroid: Vec<f32> = if dim > 0 {
            let mut sum = vec![0.0f32; dim];
            let mut valid = 0usize;
            for (_, emb) in &members {
                if emb.len() == dim {
                    for (i, v) in emb.iter().enumerate() {
                        sum[i] += v;
                    }
                    valid += 1;
                }
            }
            if valid > 0 {
                let norm_factor = valid as f32;
                let mut centroid: Vec<f32> = sum.into_iter().map(|v| v / norm_factor).collect();
                // L2-normalize for cosine consistency.
                let norm: f32 = centroid.iter().map(|x| x * x).sum::<f32>().sqrt();
                if norm > 1e-8 {
                    for v in &mut centroid {
                        *v /= norm;
                    }
                }
                centroid
            } else {
                vec![0.0f32; dim]
            }
        } else {
            Vec::new()
        };

        let summary_memory = Memory {
            id: summary_id.clone(),
            memory_type: MemoryType::Semantic,
            content: summary_content,
            agent_id: agent_id.clone(),
            session_id: None,
            importance: max_importance,
            tags: vec!["compressed".to_string()],
            metadata: Some(serde_json::json!({
                "compressed_from": cluster.len(),
                "compressed_at": now_secs,
            })),
            created_at: oldest_created_at,
            last_accessed_at: now_secs,
            access_count: 0,
            ttl_seconds: None,
            expires_at: None,
        };

        summaries.push((summary_memory, centroid));
        result.summary_ids.push(summary_id);

        // Soft-deprecate originals.
        for (mem, emb) in &members {
            let dep = Memory {
                expires_at: Some(expires_at),
                ..(*mem).clone()
            };
            result.deprecated_ids.push(dep.id.clone());
            deprecated.push((dep, (*emb).clone()));
        }
    }

    result.summaries_created = summaries.len();
    result.originals_deprecated = deprecated.len();

    (summaries, deprecated, result)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn mk_mem(id: &str, content: &str, importance: f32) -> Memory {
        Memory {
            id: id.to_string(),
            memory_type: MemoryType::Episodic,
            content: content.to_string(),
            agent_id: "agent1".to_string(),
            session_id: None,
            importance,
            tags: vec![],
            metadata: None,
            created_at: 1_000_000,
            last_accessed_at: 1_000_000,
            access_count: 0,
            ttl_seconds: None,
            expires_at: None,
        }
    }

    fn near_vec(base: &[f32], noise: f32) -> Vec<f32> {
        let mut v: Vec<f32> = base
            .iter()
            .enumerate()
            .map(|(i, x)| x + if i == 0 { noise } else { 0.0 })
            .collect();
        let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        for x in &mut v {
            *x /= norm;
        }
        v
    }

    #[test]
    fn test_compress_empty() {
        let (summaries, deprecated, result) = compress_memories(&[]);
        assert!(summaries.is_empty());
        assert!(deprecated.is_empty());
        assert_eq!(result.clusters_found, 0);
    }

    #[test]
    fn test_compress_high_importance_skipped() {
        // High-importance memories (> 0.6) should not be compressed.
        let base = vec![1.0f32, 0.0, 0.0, 0.0];
        let pairs = vec![
            (mk_mem("a", "content a", 0.9), near_vec(&base, 0.01)),
            (mk_mem("b", "content b", 0.8), near_vec(&base, 0.02)),
        ];
        let (summaries, _, result) = compress_memories(&pairs);
        assert_eq!(result.memories_scanned, 0);
        assert!(summaries.is_empty());
    }

    #[test]
    fn test_compress_two_similar_low_importance() {
        let base = vec![1.0f32, 0.0, 0.0, 0.0];
        let pairs = vec![
            (
                mk_mem("a", "The API latency is high", 0.4),
                near_vec(&base, 0.01),
            ),
            (
                mk_mem("b", "API response times are slow", 0.3),
                near_vec(&base, 0.02),
            ),
        ];
        let (summaries, deprecated, result) = compress_memories(&pairs);
        assert_eq!(result.clusters_found, 1);
        assert_eq!(result.summaries_created, 1);
        assert_eq!(result.originals_deprecated, 2);
        assert_eq!(summaries.len(), 1);
        assert_eq!(deprecated.len(), 2);
        // Summary content should be a concatenation of both.
        assert!(summaries[0].0.content.contains("API latency"));
        assert!(summaries[0].0.content.contains("API response"));
        // Deprecated originals should have expires_at set.
        for (m, _) in &deprecated {
            assert!(m.expires_at.is_some());
        }
    }

    #[test]
    fn test_compress_orthogonal_no_cluster() {
        let pairs = vec![
            (
                mk_mem("a", "vector search", 0.4),
                vec![1.0f32, 0.0, 0.0, 0.0],
            ),
            (
                mk_mem("b", "graph traversal", 0.3),
                vec![0.0f32, 1.0, 0.0, 0.0],
            ),
        ];
        let (summaries, _, result) = compress_memories(&pairs);
        assert_eq!(result.clusters_found, 0);
        assert!(summaries.is_empty());
    }

    #[test]
    fn test_detect_near_duplicate() {
        use crate::consolidate::detect_near_duplicate;

        let candidates = vec![
            ("mem_1".to_string(), vec![1.0f32, 0.0, 0.0]),
            ("mem_2".to_string(), vec![0.0f32, 1.0, 0.0]),
        ];

        // Near-identical to mem_1
        let dup = vec![0.9999f32, 0.01, 0.0];
        let norm: f32 = dup.iter().map(|x| x * x).sum::<f32>().sqrt();
        let dup: Vec<f32> = dup.into_iter().map(|x| x / norm).collect();

        let found = detect_near_duplicate(&candidates, &dup, 0.95);
        assert_eq!(found, Some("mem_1".to_string()));

        // Orthogonal — no match
        let ortho = vec![0.0f32, 0.0, 1.0];
        assert!(detect_near_duplicate(&candidates, &ortho, 0.95).is_none());
    }
}