engine/

autopilot.rs

//! Auto-Pilot Knowledge Management for Dakera AI Agent Memory Platform.
//!
//! Background tasks that automatically deduplicate and consolidate memories
//! to keep the memory space clean and efficient.
//!
//! - **Auto-dedup** (every N hours): finds near-duplicate memories above a
//!   similarity threshold and keeps the one with the highest retention score.
//! - **Auto-consolidation** (every N hours): merges clusters of 3+ low-importance
//!   memories sharing 2+ tags into single Semantic memories.

use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use common::{Memory, MemoryType};
use storage::VectorStorage;
use tokio::sync::RwLock;
use tracing;

/// Configuration for the auto-pilot background tasks.
#[derive(Clone)]
pub struct AutoPilotConfig {
    /// Master switch for all auto-pilot tasks.
    pub enabled: bool,
    /// Similarity threshold for deduplication (0.0–1.0).
    pub dedup_threshold: f32,
    /// How often to run deduplication (in hours).
    pub dedup_interval_hours: u64,
    /// How often to run consolidation (in hours).
    pub consolidation_interval_hours: u64,
}

impl Default for AutoPilotConfig {
    fn default() -> Self {
        Self {
            enabled: true,
            dedup_threshold: 0.93,
            dedup_interval_hours: 6,
            consolidation_interval_hours: 12,
        }
    }
}

impl AutoPilotConfig {
    /// Load configuration from environment variables.
    pub fn from_env() -> Self {
        let enabled: bool = std::env::var("DAKERA_AUTOPILOT_ENABLED")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(true);

        let dedup_threshold: f32 = std::env::var("DAKERA_AUTOPILOT_DEDUP_THRESHOLD")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(0.93);

        let dedup_interval_hours: u64 = std::env::var("DAKERA_AUTOPILOT_DEDUP_INTERVAL_HOURS")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(6);

        let consolidation_interval_hours: u64 =
            std::env::var("DAKERA_AUTOPILOT_CONSOLIDATION_INTERVAL_HOURS")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(12);

        Self {
            enabled,
            dedup_threshold,
            dedup_interval_hours,
            consolidation_interval_hours,
        }
    }
}

/// Result of a deduplication cycle.
#[derive(Debug, Default)]
pub struct DedupResult {
    pub namespaces_processed: usize,
    pub memories_scanned: usize,
    pub duplicates_removed: usize,
}

/// Result of a consolidation cycle.
#[derive(Debug, Default)]
pub struct ConsolidationResult {
    pub namespaces_processed: usize,
    pub memories_scanned: usize,
    pub clusters_merged: usize,
    pub memories_consolidated: usize,
}

/// Auto-Pilot engine that runs deduplication and consolidation as background tasks.
pub struct AutoPilotEngine {
    pub config: AutoPilotConfig,
}

impl AutoPilotEngine {
    pub fn new(config: AutoPilotConfig) -> Self {
        Self { config }
    }

    /// Compute cosine similarity between two embedding vectors.
    fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
        if a.len() != b.len() || a.is_empty() {
            return 0.0;
        }
        let mut dot = 0.0_f64;
        let mut norm_a = 0.0_f64;
        let mut norm_b = 0.0_f64;
        for (x, y) in a.iter().zip(b.iter()) {
            let xd = *x as f64;
            let yd = *y as f64;
            dot += xd * yd;
            norm_a += xd * xd;
            norm_b += yd * yd;
        }
        let denom = norm_a.sqrt() * norm_b.sqrt();
        if denom == 0.0 {
            0.0
        } else {
            (dot / denom) as f32
        }
    }

    /// Retention score: higher = keep this memory over duplicates.
    fn retention_score(memory: &Memory) -> f64 {
        memory.importance as f64 + memory.access_count as f64 * 0.01
    }

    /// Run deduplication across all agent namespaces.
    ///
    /// For each `_dakera_agent_*` namespace, compares all memory pairs using
    /// cosine similarity on their stored embeddings. When a pair exceeds the
    /// threshold, the memory with the lower retention score is deleted.
    pub async fn run_dedup(&self, storage: &Arc<dyn VectorStorage>) -> DedupResult {
        let mut result = DedupResult::default();

        let namespaces = match storage.list_namespaces().await {
            Ok(ns) => ns,
            Err(e) => {
                tracing::error!(error = %e, "Auto-dedup: failed to list namespaces");
                return result;
            }
        };

        for namespace in namespaces {
            if !namespace.starts_with("_dakera_agent_") {
                continue;
            }
            result.namespaces_processed += 1;

            let vectors = match storage.get_all(&namespace).await {
                Ok(v) => v,
                Err(e) => {
                    tracing::warn!(
                        namespace = %namespace,
                        error = %e,
                        "Auto-dedup: failed to get vectors"
                    );
                    continue;
                }
            };

            // Parse memories paired with their stored embeddings
            let items: Vec<(Memory, &[f32])> = vectors
                .iter()
                .filter_map(|v| {
                    let mem = Memory::from_vector(v)?;
                    if v.values.is_empty() {
                        return None;
                    }
                    Some((mem, v.values.as_slice()))
                })
                .collect();

            result.memories_scanned += items.len();

            // Pairwise comparison — O(n²) but runs infrequently on bounded namespaces
            let mut to_delete: HashSet<String> = HashSet::new();

            for i in 0..items.len() {
                if to_delete.contains(&items[i].0.id) {
                    continue;
                }
                for j in (i + 1)..items.len() {
                    if to_delete.contains(&items[j].0.id) {
                        continue;
                    }
                    let sim = Self::cosine_similarity(items[i].1, items[j].1);
                    if sim >= self.config.dedup_threshold {
                        // Keep the one with the higher retention score
                        if Self::retention_score(&items[i].0) >= Self::retention_score(&items[j].0)
                        {
                            to_delete.insert(items[j].0.id.clone());
                        } else {
                            to_delete.insert(items[i].0.id.clone());
                            break; // i is deleted, skip remaining j comparisons
                        }
                    }
                }
            }

            if !to_delete.is_empty() {
                let ids: Vec<String> = to_delete.into_iter().collect();
                result.duplicates_removed += ids.len();
                if let Err(e) = storage.delete(&namespace, &ids).await {
                    tracing::warn!(
                        namespace = %namespace,
                        count = ids.len(),
                        error = %e,
                        "Auto-dedup: failed to delete duplicates"
                    );
                }
            }
        }

        tracing::info!(
            namespaces = result.namespaces_processed,
            scanned = result.memories_scanned,
            removed = result.duplicates_removed,
            "Auto-dedup cycle completed"
        );

        result
    }

    /// Run consolidation across all agent namespaces.
    ///
    /// Finds clusters of 3+ memories sharing 2+ tags where ALL members have
    /// importance < 0.3, then merges each cluster into a single Semantic memory
    /// with combined content, union of tags, and max importance.
    pub async fn run_consolidation(&self, storage: &Arc<dyn VectorStorage>) -> ConsolidationResult {
        let mut result = ConsolidationResult::default();

        let namespaces = match storage.list_namespaces().await {
            Ok(ns) => ns,
            Err(e) => {
                tracing::error!(error = %e, "Auto-consolidation: failed to list namespaces");
                return result;
            }
        };

        for namespace in namespaces {
            if !namespace.starts_with("_dakera_agent_") {
                continue;
            }
            result.namespaces_processed += 1;

            let vectors = match storage.get_all(&namespace).await {
                Ok(v) => v,
                Err(e) => {
                    tracing::warn!(
                        namespace = %namespace,
                        error = %e,
                        "Auto-consolidation: failed to get vectors"
                    );
                    continue;
                }
            };

            // Only consider low-importance memories with embeddings and tags
            let items: Vec<(Memory, Vec<f32>)> = vectors
                .iter()
                .filter_map(|v| {
                    let mem = Memory::from_vector(v)?;
                    if mem.importance >= 0.3 || v.values.is_empty() || mem.tags.is_empty() {
                        return None;
                    }
                    Some((mem, v.values.clone()))
                })
                .collect();

            result.memories_scanned += items.len();

            if items.len() < 3 {
                continue;
            }

            // Build tag → memory-index mapping
            let mut tag_to_indices: HashMap<&str, Vec<usize>> = HashMap::new();
            for (i, (mem, _)) in items.iter().enumerate() {
                for tag in &mem.tags {
                    tag_to_indices.entry(tag.as_str()).or_default().push(i);
                }
            }

            // Count shared tags between each pair
            let mut pair_shared_tags: HashMap<(usize, usize), usize> = HashMap::new();
            for indices in tag_to_indices.values() {
                for ai in 0..indices.len() {
                    for bi in (ai + 1)..indices.len() {
                        let key = (indices[ai], indices[bi]);
                        *pair_shared_tags.entry(key).or_default() += 1;
                    }
                }
            }

            // Build adjacency graph of memories sharing 2+ tags
            let mut adjacency: HashMap<usize, HashSet<usize>> = HashMap::new();
            for (&(a, b), &count) in &pair_shared_tags {
                if count >= 2 {
                    adjacency.entry(a).or_default().insert(b);
                    adjacency.entry(b).or_default().insert(a);
                }
            }

            // Find connected components (clusters of size >= 3)
            let mut visited: HashSet<usize> = HashSet::new();
            let mut clusters: Vec<Vec<usize>> = Vec::new();

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

            // Merge each qualifying cluster
            for (ci, cluster) in clusters.iter().enumerate() {
                let memories: Vec<&Memory> = cluster.iter().map(|&i| &items[i].0).collect();
                let embeddings: Vec<&Vec<f32>> = cluster.iter().map(|&i| &items[i].1).collect();

                let max_importance = memories
                    .iter()
                    .map(|m| m.importance)
                    .fold(0.0_f32, f32::max);

                // Union all tags
                let mut all_tags: Vec<String> =
                    memories.iter().flat_map(|m| m.tags.clone()).collect();
                all_tags.sort();
                all_tags.dedup();

                // Combine content
                let combined_content: String = memories
                    .iter()
                    .map(|m| m.content.as_str())
                    .collect::<Vec<_>>()
                    .join("\n---\n");

                // Average embeddings for the merged vector
                let dim = embeddings[0].len();
                let mut avg_embedding = vec![0.0_f32; dim];
                for emb in &embeddings {
                    for (i, v) in emb.iter().enumerate() {
                        avg_embedding[i] += v;
                    }
                }
                let count = embeddings.len() as f32;
                for v in &mut avg_embedding {
                    *v /= count;
                }

                let now = std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_nanos();

                let agent_id = memories[0].agent_id.clone();
                let merged_id = format!("mem_consolidated_{:x}_{}", now, ci);

                let merged_memory = Memory {
                    id: merged_id,
                    memory_type: MemoryType::Semantic,
                    content: combined_content,
                    agent_id,
                    session_id: None,
                    importance: max_importance,
                    tags: all_tags,
                    metadata: None,
                    created_at: (now / 1_000_000_000) as u64,
                    last_accessed_at: (now / 1_000_000_000) as u64,
                    access_count: 0,
                    ttl_seconds: None,
                    expires_at: None,
                };

                let merged_vector = merged_memory.to_vector(avg_embedding);

                // Delete originals, then insert merged
                let ids_to_delete: Vec<String> = memories.iter().map(|m| m.id.clone()).collect();

                if let Err(e) = storage.delete(&namespace, &ids_to_delete).await {
                    tracing::warn!(
                        namespace = %namespace,
                        error = %e,
                        "Auto-consolidation: failed to delete originals"
                    );
                    continue;
                }

                if let Err(e) = storage.upsert(&namespace, vec![merged_vector]).await {
                    tracing::warn!(
                        namespace = %namespace,
                        error = %e,
                        "Auto-consolidation: failed to insert merged memory"
                    );
                    continue;
                }

                result.clusters_merged += 1;
                result.memories_consolidated += ids_to_delete.len();
            }
        }

        tracing::info!(
            namespaces = result.namespaces_processed,
            scanned = result.memories_scanned,
            clusters = result.clusters_merged,
            consolidated = result.memories_consolidated,
            "Auto-consolidation cycle completed"
        );

        result
    }

    /// Spawn the auto-pilot as two background tokio tasks (dedup + consolidation).
    ///
    /// Takes a shared `Arc<RwLock<AutoPilotConfig>>` so that config changes made at
    /// runtime via the PILOT-2 endpoint (`PUT /admin/autopilot/config`) take effect
    /// without a server restart. Each loop iteration re-reads the config before
    /// sleeping, so interval and threshold changes apply on the next cycle.
    pub fn spawn(
        config: Arc<RwLock<AutoPilotConfig>>,
        storage: Arc<dyn VectorStorage>,
        metrics: Arc<crate::decay::BackgroundMetrics>,
    ) -> (tokio::task::JoinHandle<()>, tokio::task::JoinHandle<()>) {
        let storage_dedup = storage.clone();
        let metrics_dedup = metrics.clone();
        let config_dedup = config.clone();

        let dedup_handle = tokio::spawn(async move {
            loop {
                let (enabled, dedup_threshold, interval_hours) = {
                    let cfg = config_dedup.read().await;
                    (cfg.enabled, cfg.dedup_threshold, cfg.dedup_interval_hours)
                };

                if !enabled {
                    // Retry every 5 minutes while disabled
                    tokio::time::sleep(std::time::Duration::from_secs(300)).await;
                    continue;
                }

                tokio::time::sleep(std::time::Duration::from_secs(interval_hours * 3600)).await;

                // Re-check enabled after sleep — may have been disabled during the interval
                if !config_dedup.read().await.enabled {
                    continue;
                }

                let engine = AutoPilotEngine::new(AutoPilotConfig {
                    enabled: true,
                    dedup_threshold,
                    ..Default::default()
                });
                let result = engine.run_dedup(&storage_dedup).await;
                metrics_dedup.record_dedup(
                    result.namespaces_processed,
                    result.memories_scanned,
                    result.duplicates_removed,
                );
            }
        });

        let consolidation_handle = tokio::spawn(async move {
            loop {
                let (enabled, interval_hours) = {
                    let cfg = config.read().await;
                    (cfg.enabled, cfg.consolidation_interval_hours)
                };

                if !enabled {
                    tokio::time::sleep(std::time::Duration::from_secs(300)).await;
                    continue;
                }

                tokio::time::sleep(std::time::Duration::from_secs(interval_hours * 3600)).await;

                if !config.read().await.enabled {
                    continue;
                }

                let engine = AutoPilotEngine::new(AutoPilotConfig::default());
                let result = engine.run_consolidation(&storage).await;
                metrics.record_consolidation(
                    result.namespaces_processed,
                    result.memories_scanned,
                    result.clusters_merged,
                    result.memories_consolidated,
                );
            }
        });

        (dedup_handle, consolidation_handle)
    }
}

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

    #[test]
    fn test_cosine_similarity_identical() {
        let a = vec![1.0, 0.0, 0.0];
        let b = vec![1.0, 0.0, 0.0];
        let sim = AutoPilotEngine::cosine_similarity(&a, &b);
        assert!((sim - 1.0).abs() < 0.001);
    }

    #[test]
    fn test_cosine_similarity_orthogonal() {
        let a = vec![1.0, 0.0, 0.0];
        let b = vec![0.0, 1.0, 0.0];
        let sim = AutoPilotEngine::cosine_similarity(&a, &b);
        assert!(sim.abs() < 0.001);
    }

    #[test]
    fn test_cosine_similarity_opposite() {
        let a = vec![1.0, 0.0];
        let b = vec![-1.0, 0.0];
        let sim = AutoPilotEngine::cosine_similarity(&a, &b);
        assert!((sim - (-1.0)).abs() < 0.001);
    }

    #[test]
    fn test_cosine_similarity_empty() {
        let sim = AutoPilotEngine::cosine_similarity(&[], &[]);
        assert!(sim.abs() < 0.001);
    }

    #[test]
    fn test_retention_score() {
        let mut mem = Memory {
            id: "test".to_string(),
            memory_type: MemoryType::Episodic,
            content: "test".to_string(),
            agent_id: "agent".to_string(),
            session_id: None,
            importance: 0.5,
            tags: vec![],
            metadata: None,
            created_at: 0,
            last_accessed_at: 0,
            access_count: 10,
            ttl_seconds: None,
            expires_at: None,
        };
        let score_a = AutoPilotEngine::retention_score(&mem);

        mem.importance = 0.8;
        mem.access_count = 0;
        let score_b = AutoPilotEngine::retention_score(&mem);

        // 0.5 + 10*0.01 = 0.6  vs  0.8 + 0 = 0.8
        assert!((score_a - 0.6).abs() < 0.001);
        assert!((score_b - 0.8).abs() < 0.001);
    }

    #[test]
    fn test_config_defaults() {
        let config = AutoPilotConfig::default();
        assert!(config.enabled);
        assert!((config.dedup_threshold - 0.93).abs() < 0.001);
        assert_eq!(config.dedup_interval_hours, 6);
        assert_eq!(config.consolidation_interval_hours, 12);
    }
}