Skip to main content

agent_io/memory/
entry.rs

1//! Memory entry types
2
3use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6use std::collections::HashMap;
7
8/// Memory type classification
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
10#[serde(rename_all = "snake_case")]
11pub enum MemoryType {
12    /// Recent conversation (short-term)
13    #[default]
14    ShortTerm,
15    /// Persistent knowledge (long-term)
16    LongTerm,
17    /// Specific events/experiences
18    Episodic,
19    /// Facts and concepts
20    Semantic,
21}
22
23/// Memory entry
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct MemoryEntry {
26    /// Unique identifier
27    pub id: String,
28    /// Memory content
29    pub content: String,
30    /// Embedding vector (for similarity search)
31    #[serde(skip_serializing_if = "Option::is_none")]
32    pub embedding: Option<Vec<f32>>,
33    /// Memory type
34    #[serde(default)]
35    pub memory_type: MemoryType,
36    /// Additional metadata
37    #[serde(default)]
38    pub metadata: HashMap<String, Value>,
39    /// Creation timestamp
40    pub created_at: DateTime<Utc>,
41    /// Last access timestamp
42    #[serde(skip_serializing_if = "Option::is_none")]
43    pub last_accessed: Option<DateTime<Utc>>,
44    /// Importance score (0.0 - 1.0)
45    #[serde(default = "default_importance")]
46    pub importance: f32,
47    /// Access count (for recency weighting)
48    #[serde(default)]
49    pub access_count: u32,
50}
51
52fn default_importance() -> f32 {
53    0.5
54}
55
56impl MemoryEntry {
57    /// Create a new memory entry
58    pub fn new(content: impl Into<String>) -> Self {
59        Self {
60            id: uuid::Uuid::new_v4().to_string(),
61            content: content.into(),
62            embedding: None,
63            memory_type: MemoryType::default(),
64            metadata: HashMap::new(),
65            created_at: Utc::now(),
66            last_accessed: None,
67            importance: 0.5,
68            access_count: 0,
69        }
70    }
71
72    /// Set memory type
73    pub fn with_type(mut self, memory_type: MemoryType) -> Self {
74        self.memory_type = memory_type;
75        self
76    }
77
78    /// Set embedding
79    pub fn with_embedding(mut self, embedding: Vec<f32>) -> Self {
80        self.embedding = Some(embedding);
81        self
82    }
83
84    /// Set importance
85    pub fn with_importance(mut self, importance: f32) -> Self {
86        self.importance = importance.clamp(0.0, 1.0);
87        self
88    }
89
90    /// Add metadata
91    pub fn with_metadata(mut self, key: impl Into<String>, value: Value) -> Self {
92        self.metadata.insert(key.into(), value);
93        self
94    }
95
96    /// Record access (updates access count and timestamp)
97    pub fn record_access(&mut self) {
98        self.access_count += 1;
99        self.last_accessed = Some(Utc::now());
100    }
101
102    /// Calculate relevance score based on importance, recency, and access frequency
103    pub fn relevance_score(&self) -> f32 {
104        let age_hours = (Utc::now() - self.created_at).num_hours() as f32;
105        let recency_factor = (-age_hours / 24.0 / 7.0).exp(); // Decay over a week
106
107        let access_factor = 1.0 + (self.access_count as f32).ln().max(0.0) * 0.1;
108
109        self.importance * recency_factor * access_factor
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116
117    #[test]
118    fn test_memory_entry_creation() {
119        let entry = MemoryEntry::new("Test memory content");
120        assert_eq!(entry.content, "Test memory content");
121        assert_eq!(entry.memory_type, MemoryType::ShortTerm);
122        assert_eq!(entry.importance, 0.5);
123    }
124
125    #[test]
126    fn test_memory_entry_builder() {
127        let entry = MemoryEntry::new("Test")
128            .with_type(MemoryType::LongTerm)
129            .with_importance(0.9)
130            .with_metadata("source", Value::String("user".to_string()));
131
132        assert_eq!(entry.memory_type, MemoryType::LongTerm);
133        assert_eq!(entry.importance, 0.9);
134        assert!(entry.metadata.contains_key("source"));
135    }
136
137    #[test]
138    fn test_record_access() {
139        let mut entry = MemoryEntry::new("Test");
140        assert_eq!(entry.access_count, 0);
141
142        entry.record_access();
143        assert_eq!(entry.access_count, 1);
144        assert!(entry.last_accessed.is_some());
145    }
146}