1use chrono::{DateTime, Utc};
2use regex::Regex;
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5use std::sync::LazyLock;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
8#[serde(rename_all = "lowercase")]
9pub enum MemoryCategory {
10 Preference,
11 Fact,
12 Decision,
13 Entity,
14 Other,
15}
16
17static PREFERENCE_RE: LazyLock<Regex> = LazyLock::new(|| {
18 Regex::new(r"(?i)\b(prefer|like|want|favorite|rather|love|enjoy|wish|choose to|opt for)\b")
19 .unwrap()
20});
21static DECISION_RE: LazyLock<Regex> = LazyLock::new(|| {
22 Regex::new(
23 r"(?i)\b(decided|chose|going with|will use|picked|settled on|switched to|migrated to)\b",
24 )
25 .unwrap()
26});
27static FACT_RE: LazyLock<Regex> = LazyLock::new(|| {
28 Regex::new(r"(?i)\b(is|are|was|has|have|knows? that|works? (on|at|with)|uses?|built with)\b")
29 .unwrap()
30});
31static ENTITY_RE: LazyLock<Regex> =
32 LazyLock::new(|| Regex::new(r"\b[A-Z][a-z]+(?:\s+[A-Z][a-z]+)+\b").unwrap());
33
34impl MemoryCategory {
35 pub fn detect(text: &str) -> Self {
36 if PREFERENCE_RE.is_match(text) {
37 Self::Preference
38 } else if DECISION_RE.is_match(text) {
39 Self::Decision
40 } else if ENTITY_RE.is_match(text) {
41 Self::Entity
42 } else if FACT_RE.is_match(text) {
43 Self::Fact
44 } else {
45 Self::Other
46 }
47 }
48
49 pub fn as_str(&self) -> &'static str {
50 match self {
51 Self::Preference => "preference",
52 Self::Fact => "fact",
53 Self::Decision => "decision",
54 Self::Entity => "entity",
55 Self::Other => "other",
56 }
57 }
58
59 pub fn from_str_lossy(s: &str) -> Self {
60 match s.to_lowercase().as_str() {
61 "preference" => Self::Preference,
62 "fact" => Self::Fact,
63 "decision" => Self::Decision,
64 "entity" => Self::Entity,
65 _ => Self::Other,
66 }
67 }
68}
69
70impl std::fmt::Display for MemoryCategory {
71 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
72 f.write_str(self.as_str())
73 }
74}
75
76#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
77#[serde(rename_all = "lowercase")]
78pub enum MemorySource {
79 Manual,
80 Conversation,
81 Session,
82}
83
84impl MemorySource {
85 pub fn as_str(&self) -> &'static str {
86 match self {
87 Self::Manual => "manual",
88 Self::Conversation => "conversation",
89 Self::Session => "session",
90 }
91 }
92
93 pub fn from_str_lossy(s: &str) -> Self {
94 match s.to_lowercase().as_str() {
95 "conversation" => Self::Conversation,
96 "session" => Self::Session,
97 _ => Self::Manual,
98 }
99 }
100}
101
102#[derive(Debug, Clone, Serialize, Deserialize)]
103pub struct MemoryEntry {
104 pub id: String,
105 pub content: String,
106 pub source: MemorySource,
107 pub category: MemoryCategory,
108 pub container_tag: String,
109 pub metadata: HashMap<String, serde_json::Value>,
110 pub session_id: Option<String>,
111 pub created_at: DateTime<Utc>,
112 pub updated_at: DateTime<Utc>,
113}
114
115impl MemoryEntry {
116 pub fn new(
117 content: impl Into<String>,
118 source: MemorySource,
119 container_tag: impl Into<String>,
120 ) -> Self {
121 let content = content.into();
122 let category = MemoryCategory::detect(&content);
123 let now = Utc::now();
124 Self {
125 id: uuid::Uuid::new_v4().to_string(),
126 content,
127 source,
128 category,
129 container_tag: container_tag.into(),
130 metadata: HashMap::new(),
131 session_id: None,
132 created_at: now,
133 updated_at: now,
134 }
135 }
136
137 pub fn with_category(mut self, category: MemoryCategory) -> Self {
138 self.category = category;
139 self
140 }
141
142 pub fn with_session_id(mut self, session_id: impl Into<String>) -> Self {
143 self.session_id = Some(session_id.into());
144 self
145 }
146
147 pub fn with_metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
148 self.metadata.insert(key.into(), value);
149 self
150 }
151}
152
153#[derive(Debug, Clone, Serialize, Deserialize)]
154pub struct SearchResult {
155 pub entry: MemoryEntry,
156 pub vector_score: Option<f32>,
157 pub fts_score: Option<f64>,
158 pub combined_score: f64,
159 pub final_score: f64,
160}
161
162#[derive(Debug, Clone, Serialize, Deserialize)]
163pub struct UserProfileFact {
164 pub id: String,
165 pub fact_type: String,
166 pub category: String,
167 pub content: String,
168 pub created_at: DateTime<Utc>,
169 pub updated_at: DateTime<Utc>,
170}
171
172#[derive(Debug, Clone, Default, Serialize, Deserialize)]
173pub struct MemoryStats {
174 pub total_memories: usize,
175 pub total_by_category: HashMap<String, usize>,
176 pub total_by_container: HashMap<String, usize>,
177 pub db_size_bytes: u64,
178}
179
180#[cfg(test)]
181mod tests {
182 use super::*;
183
184 #[test]
185 fn detect_preference() {
186 assert_eq!(
187 MemoryCategory::detect("I prefer dark mode"),
188 MemoryCategory::Preference
189 );
190 assert_eq!(
191 MemoryCategory::detect("I like using Rust"),
192 MemoryCategory::Preference
193 );
194 assert_eq!(
195 MemoryCategory::detect("I want fast builds"),
196 MemoryCategory::Preference
197 );
198 assert_eq!(
199 MemoryCategory::detect("My favorite editor is Neovim"),
200 MemoryCategory::Preference
201 );
202 }
203
204 #[test]
205 fn detect_decision() {
206 assert_eq!(
207 MemoryCategory::detect("We decided to use SQLite"),
208 MemoryCategory::Decision
209 );
210 assert_eq!(
211 MemoryCategory::detect("I chose rusqlite over LanceDB"),
212 MemoryCategory::Decision
213 );
214 assert_eq!(
215 MemoryCategory::detect("Going with the hybrid approach"),
216 MemoryCategory::Decision
217 );
218 }
219
220 #[test]
221 fn detect_entity() {
222 assert_eq!(
223 MemoryCategory::detect("John Smith works here"),
224 MemoryCategory::Entity
225 );
226 assert_eq!(
227 MemoryCategory::detect("Junho Yeo created uira"),
228 MemoryCategory::Entity
229 );
230 }
231
232 #[test]
233 fn detect_fact() {
234 assert_eq!(
235 MemoryCategory::detect("The project is built with Rust"),
236 MemoryCategory::Fact
237 );
238 assert_eq!(
239 MemoryCategory::detect("It uses tokio for async"),
240 MemoryCategory::Fact
241 );
242 }
243
244 #[test]
245 fn detect_other() {
246 assert_eq!(MemoryCategory::detect("hello"), MemoryCategory::Other);
247 assert_eq!(MemoryCategory::detect("1234"), MemoryCategory::Other);
248 }
249
250 #[test]
251 fn memory_entry_auto_categorizes() {
252 let entry = MemoryEntry::new("I prefer dark mode", MemorySource::Manual, "default");
253 assert_eq!(entry.category, MemoryCategory::Preference);
254 assert!(!entry.id.is_empty());
255 }
256
257 #[test]
258 fn memory_entry_builder() {
259 let entry = MemoryEntry::new("test", MemorySource::Conversation, "default")
260 .with_category(MemoryCategory::Fact)
261 .with_session_id("ses_123")
262 .with_metadata("key", serde_json::json!("value"));
263 assert_eq!(entry.category, MemoryCategory::Fact);
264 assert_eq!(entry.session_id, Some("ses_123".to_string()));
265 assert_eq!(entry.metadata.get("key"), Some(&serde_json::json!("value")));
266 }
267
268 #[test]
269 fn category_roundtrip() {
270 for cat in [
271 MemoryCategory::Preference,
272 MemoryCategory::Fact,
273 MemoryCategory::Decision,
274 MemoryCategory::Entity,
275 MemoryCategory::Other,
276 ] {
277 assert_eq!(MemoryCategory::from_str_lossy(cat.as_str()), cat);
278 }
279 }
280
281 #[test]
282 fn source_roundtrip() {
283 for src in [
284 MemorySource::Manual,
285 MemorySource::Conversation,
286 MemorySource::Session,
287 ] {
288 assert_eq!(MemorySource::from_str_lossy(src.as_str()), src);
289 }
290 }
291}