1use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6use std::collections::HashMap;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
10#[serde(rename_all = "snake_case")]
11pub enum MemoryType {
12 #[default]
14 ShortTerm,
15 LongTerm,
17 Episodic,
19 Semantic,
21}
22
23#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct MemoryEntry {
26 pub id: String,
28 pub content: String,
30 #[serde(skip_serializing_if = "Option::is_none")]
32 pub embedding: Option<Vec<f32>>,
33 #[serde(default)]
35 pub memory_type: MemoryType,
36 #[serde(default)]
38 pub metadata: HashMap<String, Value>,
39 pub created_at: DateTime<Utc>,
41 #[serde(skip_serializing_if = "Option::is_none")]
43 pub last_accessed: Option<DateTime<Utc>>,
44 #[serde(default = "default_importance")]
46 pub importance: f32,
47 #[serde(default)]
49 pub access_count: u32,
50}
51
52fn default_importance() -> f32 {
53 0.5
54}
55
56impl MemoryEntry {
57 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 pub fn with_type(mut self, memory_type: MemoryType) -> Self {
74 self.memory_type = memory_type;
75 self
76 }
77
78 pub fn with_embedding(mut self, embedding: Vec<f32>) -> Self {
80 self.embedding = Some(embedding);
81 self
82 }
83
84 pub fn with_importance(mut self, importance: f32) -> Self {
86 self.importance = importance.clamp(0.0, 1.0);
87 self
88 }
89
90 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 pub fn record_access(&mut self) {
98 self.access_count += 1;
99 self.last_accessed = Some(Utc::now());
100 }
101
102 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(); 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}