rust-memory-service/src/storage/entry.rs

//! Memory Entry Definition
//!
//! Core data structure for memory entries with serialization support.

use serde::{Deserialize, Serialize};
use chrono::{DateTime, Utc};
use std::collections::HashMap;

/// Memory types as defined in CoALA paper
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[repr(u8)]
pub enum MemoryType {
    /// Specific events/interactions
    Episodic = 0,
    /// General knowledge/facts
    Semantic = 1,
    /// Learned patterns/skills
    Procedural = 2,
    /// Time-aware context
    Temporal = 3,
}

impl From<u8> for MemoryType {
    fn from(value: u8) -> Self {
        match value {
            0 => MemoryType::Episodic,
            1 => MemoryType::Semantic,
            2 => MemoryType::Procedural,
            3 => MemoryType::Temporal,
            _ => MemoryType::Episodic,
        }
    }
}

impl std::fmt::Display for MemoryType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            MemoryType::Episodic => write!(f, "episodic"),
            MemoryType::Semantic => write!(f, "semantic"),
            MemoryType::Procedural => write!(f, "procedural"),
            MemoryType::Temporal => write!(f, "temporal"),
        }
    }
}

/// A single memory entry
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryEntry {
    /// Unique identifier
    pub id: String,

    /// Memory type category
    #[serde(rename = "type")]
    pub memory_type: MemoryType,

    /// The actual content/text
    pub content: String,

    /// Vector embedding for similarity search
    #[serde(skip_serializing_if = "Option::is_none")]
    pub embedding: Option<Vec<f32>>,

    /// Additional structured metadata
    #[serde(default)]
    pub metadata: HashMap<String, serde_json::Value>,

    /// When this memory was created
    pub timestamp: DateTime<Utc>,

    /// Time-to-live in seconds (None = no expiration)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub ttl_seconds: Option<u64>,

    /// Confidence score (0.0 - 1.0)
    pub confidence: f64,

    /// User this memory belongs to
    pub user_id: String,

    /// Session this memory was created in
    pub session_id: String,

    /// Source of this memory
    pub source: String,

    /// Tags for categorization
    #[serde(default)]
    pub tags: Vec<String>,

    /// Version for optimistic concurrency
    pub version: i32,

    /// Dependencies on other memories
    #[serde(default)]
    pub dependencies: Vec<String>,
}

impl MemoryEntry {
    /// Create a new memory entry with sensible defaults
    pub fn new(
        id: String,
        memory_type: MemoryType,
        content: String,
        user_id: String,
        session_id: String,
    ) -> Self {
        Self {
            id,
            memory_type,
            content,
            embedding: None,
            metadata: HashMap::new(),
            timestamp: Utc::now(),
            ttl_seconds: None,
            confidence: 0.8,
            user_id,
            session_id,
            source: "conversation".to_string(),
            tags: Vec::new(),
            version: 1,
            dependencies: Vec::new(),
        }
    }

    /// Check if this entry has expired
    pub fn is_expired(&self) -> bool {
        if let Some(ttl) = self.ttl_seconds {
            let expiry = self.timestamp + chrono::Duration::seconds(ttl as i64);
            Utc::now() > expiry
        } else {
            false
        }
    }

    /// Calculate expiry time if TTL is set
    pub fn expires_at(&self) -> Option<DateTime<Utc>> {
        self.ttl_seconds.map(|ttl| {
            self.timestamp + chrono::Duration::seconds(ttl as i64)
        })
    }

    /// Calculate age in hours
    pub fn age_hours(&self) -> f64 {
        let duration = Utc::now() - self.timestamp;
        duration.num_seconds() as f64 / 3600.0
    }

    /// Set embedding
    pub fn with_embedding(mut self, embedding: Vec<f32>) -> Self {
        self.embedding = Some(embedding);
        self
    }

    /// Set TTL
    pub fn with_ttl(mut self, seconds: u64) -> Self {
        self.ttl_seconds = Some(seconds);
        self
    }

    /// Set confidence
    pub fn with_confidence(mut self, confidence: f64) -> Self {
        self.confidence = confidence;
        self
    }

    /// Set source
    pub fn with_source(mut self, source: String) -> Self {
        self.source = source;
        self
    }

    /// Add tags
    pub fn with_tags(mut self, tags: Vec<String>) -> Self {
        self.tags = tags;
        self
    }

    /// Add metadata
    pub fn with_metadata(mut self, key: String, value: serde_json::Value) -> Self {
        self.metadata.insert(key, value);
        self
    }

    /// Calculate similarity with another entry using embeddings
    pub fn similarity(&self, other: &MemoryEntry) -> Option<f64> {
        match (&self.embedding, &other.embedding) {
            (Some(a), Some(b)) => Some(cosine_similarity(a, b)),
            _ => None,
        }
    }
}

/// Calculate cosine similarity between two vectors
fn cosine_similarity(a: &[f32], b: &[f32]) -> f64 {
    if a.len() != b.len() {
        return 0.0;
    }

    let mut dot_product = 0.0f64;
    let mut norm_a = 0.0f64;
    let mut norm_b = 0.0f64;

    for i in 0..a.len() {
        dot_product += (a[i] as f64) * (b[i] as f64);
        norm_a += (a[i] as f64).powi(2);
        norm_b += (b[i] as f64).powi(2);
    }

    if norm_a == 0.0 || norm_b == 0.0 {
        return 0.0;
    }

    dot_product / (norm_a.sqrt() * norm_b.sqrt())
}

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

    #[test]
    fn test_entry_creation() {
        let entry = MemoryEntry::new(
            "test-1".to_string(),
            MemoryType::Episodic,
            "Test content".to_string(),
            "user-1".to_string(),
            "session-1".to_string(),
        );

        assert_eq!(entry.id, "test-1");
        assert_eq!(entry.memory_type, MemoryType::Episodic);
        assert!(!entry.is_expired());
    }

    #[test]
    fn test_cosine_similarity() {
        let a = vec![1.0, 0.0, 0.0];
        let b = vec![1.0, 0.0, 0.0];
        assert!((cosine_similarity(&a, &b) - 1.0).abs() < 0.0001);

        let c = vec![0.0, 1.0, 0.0];
        assert!(cosine_similarity(&a, &c).abs() < 0.0001);
    }
}