Skip to main content

uira_memory/
types.rs

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}