ceylon_next/memory/advanced/

semantic.rs

//! Semantic Memory - Facts and Knowledge Management
//!
//! Semantic memory stores extracted facts, knowledge, and entities from conversations.
//! It provides a knowledge graph-like structure for long-term information retention.
//!
//! # Features
//!
//! - Entity extraction and tracking
//! - Fact storage and retrieval
//! - Relationship mapping
//! - Knowledge consolidation

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

/// Types of entities that can be extracted
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum EntityType {
    Person,
    Place,
    Organization,
    Concept,
    Product,
    Event,
    Other(String),
}

/// An entity extracted from conversations
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Entity {
    /// Unique identifier
    pub id: String,
    /// Entity name
    pub name: String,
    /// Type of entity
    pub entity_type: EntityType,
    /// Attributes associated with this entity
    pub attributes: HashMap<String, String>,
    /// First mentioned timestamp
    pub first_mentioned: u64,
    /// Last mentioned timestamp
    pub last_mentioned: u64,
    /// Number of mentions
    pub mention_count: u32,
    /// Related memory IDs
    pub source_memories: Vec<String>,
}

impl Entity {
    pub fn new(name: String, entity_type: EntityType, memory_id: String) -> Self {
        let timestamp = Self::current_timestamp();
        Self {
            id: uuid::Uuid::new_v4().to_string(),
            name,
            entity_type,
            attributes: HashMap::new(),
            first_mentioned: timestamp,
            last_mentioned: timestamp,
            mention_count: 1,
            source_memories: vec![memory_id],
        }
    }

    pub fn update_mention(&mut self, memory_id: String) {
        self.last_mentioned = Self::current_timestamp();
        self.mention_count += 1;
        if !self.source_memories.contains(&memory_id) {
            self.source_memories.push(memory_id);
        }
    }

    fn current_timestamp() -> u64 {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs()
    }
}

/// A fact extracted from conversations
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Fact {
    /// Unique identifier
    pub id: String,
    /// Subject entity
    pub subject: String,
    /// Predicate (relationship/property)
    pub predicate: String,
    /// Object (value/related entity)
    pub object: String,
    /// Confidence score (0.0 to 1.0)
    pub confidence: f32,
    /// When this fact was learned
    pub learned_at: u64,
    /// Last confirmed timestamp
    pub last_confirmed: u64,
    /// Confirmation count
    pub confirmation_count: u32,
    /// Source memory IDs
    pub source_memories: Vec<String>,
}

impl Fact {
    pub fn new(subject: String, predicate: String, object: String, memory_id: String) -> Self {
        let timestamp = Self::current_timestamp();
        Self {
            id: uuid::Uuid::new_v4().to_string(),
            subject,
            predicate,
            object,
            confidence: 0.7, // Default confidence
            learned_at: timestamp,
            last_confirmed: timestamp,
            confirmation_count: 1,
            source_memories: vec![memory_id],
        }
    }

    pub fn confirm(&mut self, memory_id: String) {
        self.last_confirmed = Self::current_timestamp();
        self.confirmation_count += 1;
        // Increase confidence with each confirmation (max 1.0)
        self.confidence = (self.confidence + 0.1).min(1.0);
        if !self.source_memories.contains(&memory_id) {
            self.source_memories.push(memory_id);
        }
    }

    fn current_timestamp() -> u64 {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs()
    }
}

/// Relationship between entities
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Relationship {
    pub id: String,
    pub from_entity: String,
    pub to_entity: String,
    pub relationship_type: String,
    pub strength: f32,
    pub created_at: u64,
}

/// Semantic memory for managing facts and knowledge
pub struct SemanticMemory {
    /// Entities indexed by ID
    entities: Arc<RwLock<HashMap<String, Entity>>>,
    /// Entity name to ID mapping
    entity_names: Arc<RwLock<HashMap<String, String>>>,
    /// Facts indexed by ID
    facts: Arc<RwLock<HashMap<String, Fact>>>,
    /// Subject-based fact index
    fact_index: Arc<RwLock<HashMap<String, Vec<String>>>>,
    /// Relationships
    relationships: Arc<RwLock<HashMap<String, Relationship>>>,
}

impl SemanticMemory {
    pub fn new() -> Self {
        Self {
            entities: Arc::new(RwLock::new(HashMap::new())),
            entity_names: Arc::new(RwLock::new(HashMap::new())),
            facts: Arc::new(RwLock::new(HashMap::new())),
            fact_index: Arc::new(RwLock::new(HashMap::new())),
            relationships: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    // ============================================
    // Entity Management
    // ============================================

    /// Add or update an entity
    pub async fn add_entity(&self, name: String, entity_type: EntityType, memory_id: String) -> String {
        let mut entities = self.entities.write().await;
        let mut entity_names = self.entity_names.write().await;

        // Check if entity already exists
        if let Some(entity_id) = entity_names.get(&name) {
            if let Some(entity) = entities.get_mut(entity_id) {
                entity.update_mention(memory_id);
                return entity_id.clone();
            }
        }

        // Create new entity
        let entity = Entity::new(name.clone(), entity_type, memory_id);
        let id = entity.id.clone();
        entity_names.insert(name, id.clone());
        entities.insert(id.clone(), entity);

        id
    }

    /// Get an entity by ID
    pub async fn get_entity(&self, id: &str) -> Option<Entity> {
        let entities = self.entities.read().await;
        entities.get(id).cloned()
    }

    /// Get an entity by name
    pub async fn get_entity_by_name(&self, name: &str) -> Option<Entity> {
        let entity_names = self.entity_names.read().await;
        if let Some(id) = entity_names.get(name) {
            let entities = self.entities.read().await;
            entities.get(id).cloned()
        } else {
            None
        }
    }

    /// Get all entities of a specific type
    pub async fn get_entities_by_type(&self, entity_type: EntityType) -> Vec<Entity> {
        let entities = self.entities.read().await;
        entities
            .values()
            .filter(|e| e.entity_type == entity_type)
            .cloned()
            .collect()
    }

    /// Search entities by name (partial match)
    pub async fn search_entities(&self, query: &str) -> Vec<Entity> {
        let entities = self.entities.read().await;
        let query_lower = query.to_lowercase();
        entities
            .values()
            .filter(|e| e.name.to_lowercase().contains(&query_lower))
            .cloned()
            .collect()
    }

    // ============================================
    // Fact Management
    // ============================================

    /// Add or confirm a fact
    pub async fn add_fact(
        &self,
        subject: String,
        predicate: String,
        object: String,
        memory_id: String,
    ) -> String {
        let mut facts = self.facts.write().await;
        let mut fact_index = self.fact_index.write().await;

        // Check if similar fact exists
        if let Some(fact_ids) = fact_index.get(&subject) {
            for fact_id in fact_ids {
                if let Some(fact) = facts.get_mut(fact_id) {
                    if fact.predicate == predicate && fact.object == object {
                        fact.confirm(memory_id);
                        return fact_id.clone();
                    }
                }
            }
        }

        // Create new fact
        let fact = Fact::new(subject.clone(), predicate, object, memory_id);
        let id = fact.id.clone();

        // Add to index
        fact_index
            .entry(subject)
            .or_insert_with(Vec::new)
            .push(id.clone());

        facts.insert(id.clone(), fact);

        id
    }

    /// Get a fact by ID
    pub async fn get_fact(&self, id: &str) -> Option<Fact> {
        let facts = self.facts.read().await;
        facts.get(id).cloned()
    }

    /// Get all facts about a subject
    pub async fn get_facts_about(&self, subject: &str) -> Vec<Fact> {
        let facts = self.facts.read().await;
        let fact_index = self.fact_index.read().await;

        if let Some(fact_ids) = fact_index.get(subject) {
            fact_ids
                .iter()
                .filter_map(|id| facts.get(id).cloned())
                .collect()
        } else {
            Vec::new()
        }
    }

    /// Get all facts with a specific predicate
    pub async fn get_facts_by_predicate(&self, predicate: &str) -> Vec<Fact> {
        let facts = self.facts.read().await;
        facts
            .values()
            .filter(|f| f.predicate == predicate)
            .cloned()
            .collect()
    }

    /// Get high-confidence facts (confidence >= threshold)
    pub async fn get_high_confidence_facts(&self, threshold: f32) -> Vec<Fact> {
        let facts = self.facts.read().await;
        facts
            .values()
            .filter(|f| f.confidence >= threshold)
            .cloned()
            .collect()
    }

    // ============================================
    // Relationship Management
    // ============================================

    /// Add a relationship between entities
    pub async fn add_relationship(
        &self,
        from_entity: String,
        to_entity: String,
        relationship_type: String,
    ) -> String {
        let mut relationships = self.relationships.write().await;

        let relationship = Relationship {
            id: uuid::Uuid::new_v4().to_string(),
            from_entity,
            to_entity,
            relationship_type,
            strength: 1.0,
            created_at: Self::current_timestamp(),
        };

        let id = relationship.id.clone();
        relationships.insert(id.clone(), relationship);
        id
    }

    /// Get relationships for an entity
    pub async fn get_relationships(&self, entity_id: &str) -> Vec<Relationship> {
        let relationships = self.relationships.read().await;
        relationships
            .values()
            .filter(|r| r.from_entity == entity_id || r.to_entity == entity_id)
            .cloned()
            .collect()
    }

    // ============================================
    // Knowledge Queries
    // ============================================

    /// Get a knowledge summary for an entity
    pub async fn get_knowledge_summary(&self, entity_name: &str) -> Option<KnowledgeSummary> {
        let entity = self.get_entity_by_name(entity_name).await?;
        let facts = self.get_facts_about(entity_name).await;
        let relationships = self.get_relationships(&entity.id).await;

        Some(KnowledgeSummary {
            entity: entity.clone(),
            fact_count: facts.len(),
            relationship_count: relationships.len(),
            high_confidence_facts: facts
                .iter()
                .filter(|f| f.confidence >= 0.8)
                .cloned()
                .collect(),
        })
    }

    /// Get all entity names
    pub async fn get_all_entity_names(&self) -> Vec<String> {
        let entity_names = self.entity_names.read().await;
        entity_names.keys().cloned().collect()
    }

    /// Get statistics
    pub async fn get_statistics(&self) -> SemanticMemoryStats {
        let entities = self.entities.read().await;
        let facts = self.facts.read().await;
        let relationships = self.relationships.read().await;

        SemanticMemoryStats {
            total_entities: entities.len(),
            total_facts: facts.len(),
            total_relationships: relationships.len(),
            high_confidence_facts: facts.values().filter(|f| f.confidence >= 0.8).count(),
        }
    }

    fn current_timestamp() -> u64 {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs()
    }
}

impl Default for SemanticMemory {
    fn default() -> Self {
        Self::new()
    }
}

/// Summary of knowledge about an entity
#[derive(Debug, Clone)]
pub struct KnowledgeSummary {
    pub entity: Entity,
    pub fact_count: usize,
    pub relationship_count: usize,
    pub high_confidence_facts: Vec<Fact>,
}

/// Statistics about semantic memory
#[derive(Debug, Clone)]
pub struct SemanticMemoryStats {
    pub total_entities: usize,
    pub total_facts: usize,
    pub total_relationships: usize,
    pub high_confidence_facts: usize,
}

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

    #[tokio::test]
    async fn test_semantic_memory_entities() {
        let sm = SemanticMemory::new();

        let id1 = sm
            .add_entity("Alice".to_string(), EntityType::Person, "mem-1".to_string())
            .await;
        let id2 = sm
            .add_entity("Alice".to_string(), EntityType::Person, "mem-2".to_string())
            .await;

        // Should return same ID for same entity
        assert_eq!(id1, id2);

        let entity = sm.get_entity_by_name("Alice").await.unwrap();
        assert_eq!(entity.mention_count, 2);
    }

    #[tokio::test]
    async fn test_semantic_memory_facts() {
        let sm = SemanticMemory::new();

        let id1 = sm
            .add_fact(
                "Alice".to_string(),
                "likes".to_string(),
                "programming".to_string(),
                "mem-1".to_string(),
            )
            .await;

        let id2 = sm
            .add_fact(
                "Alice".to_string(),
                "likes".to_string(),
                "programming".to_string(),
                "mem-2".to_string(),
            )
            .await;

        // Should return same ID for same fact
        assert_eq!(id1, id2);

        let fact = sm.get_fact(&id1).await.unwrap();
        assert_eq!(fact.confirmation_count, 2);
        assert!(fact.confidence > 0.7); // Should increase with confirmation
    }

    #[tokio::test]
    async fn test_semantic_memory_relationships() {
        let sm = SemanticMemory::new();

        let entity1 = sm
            .add_entity("Alice".to_string(), EntityType::Person, "mem-1".to_string())
            .await;
        let entity2 = sm
            .add_entity("Bob".to_string(), EntityType::Person, "mem-1".to_string())
            .await;

        sm.add_relationship(entity1.clone(), entity2.clone(), "friend".to_string())
            .await;

        let rels = sm.get_relationships(&entity1).await;
        assert_eq!(rels.len(), 1);
        assert_eq!(rels[0].relationship_type, "friend");
    }
}