rag_plusplus_core/store/
traits.rs

1//! Store Traits
2//!
3//! Defines the core abstraction for record storage.
4
5use crate::error::Result;
6use crate::types::{MemoryRecord, RecordId};
7use parking_lot::RwLock;
8use std::fmt::Debug;
9use std::sync::Arc;
10
11/// Core trait for record storage.
12///
13/// All storage implementations must implement this trait.
14/// The trait enforces INV-001 (record immutability) by only allowing
15/// stats updates after insertion.
16pub trait RecordStore: Send + Sync + Debug {
17    /// Insert a new record.
18    ///
19    /// # Errors
20    ///
21    /// Returns error if record with same ID already exists.
22    fn insert(&mut self, record: MemoryRecord) -> Result<RecordId>;
23
24    /// Insert multiple records in batch.
25    ///
26    /// Default implementation calls `insert` repeatedly.
27    fn insert_batch(&mut self, records: Vec<MemoryRecord>) -> Result<Vec<RecordId>> {
28        let mut ids = Vec::with_capacity(records.len());
29        for record in records {
30            ids.push(self.insert(record)?);
31        }
32        Ok(ids)
33    }
34
35    /// Get a record by ID.
36    fn get(&self, id: &RecordId) -> Option<MemoryRecord>;
37
38    /// Get multiple records by ID.
39    fn get_batch(&self, ids: &[RecordId]) -> Vec<Option<MemoryRecord>> {
40        ids.iter().map(|id| self.get(id)).collect()
41    }
42
43    /// Check if a record exists.
44    fn contains(&self, id: &RecordId) -> bool;
45
46    /// Update the outcome statistics for a record.
47    ///
48    /// This is the ONLY mutable operation allowed on a record after insertion
49    /// (enforces INV-001).
50    ///
51    /// # Arguments
52    ///
53    /// * `id` - Record ID
54    /// * `outcome` - New outcome value to incorporate
55    fn update_stats(&mut self, id: &RecordId, outcome: f64) -> Result<()>;
56
57    /// Remove a record (marks as deleted, may not physically remove).
58    ///
59    /// # Returns
60    ///
61    /// `true` if record was found and removed.
62    fn remove(&mut self, id: &RecordId) -> Result<bool>;
63
64    /// Number of records in the store.
65    fn len(&self) -> usize;
66
67    /// Whether the store is empty.
68    fn is_empty(&self) -> bool {
69        self.len() == 0
70    }
71
72    /// Clear all records.
73    fn clear(&mut self);
74
75    /// Get all record IDs.
76    fn ids(&self) -> Vec<RecordId>;
77
78    /// Get memory usage estimate in bytes.
79    fn memory_usage(&self) -> usize;
80}
81
82/// Thread-safe shared store using Arc<RwLock>.
83pub type SharedStore<S> = Arc<RwLock<S>>;
84
85/// Create a shared store from a store instance.
86#[allow(dead_code)]
87pub fn shared_store<S: RecordStore>(store: S) -> SharedStore<S> {
88    Arc::new(RwLock::new(store))
89}
90
91#[cfg(test)]
92pub mod tests {
93    use super::*;
94    use crate::OutcomeStats;
95
96    // Trait tests that any implementation should pass
97    pub fn test_basic_crud<S: RecordStore>(store: &mut S) {
98        use crate::types::RecordStatus;
99
100        // Insert
101        let record = MemoryRecord {
102            id: "test-1".into(),
103            embedding: vec![1.0, 2.0, 3.0],
104            context: "test context".into(),
105            outcome: 0.8,
106            metadata: Default::default(),
107            created_at: 0,
108            status: RecordStatus::Active,
109            stats: OutcomeStats::new(1),
110        };
111
112        let id = store.insert(record.clone()).unwrap();
113        assert_eq!(id.as_str(), "test-1");
114
115        // Get
116        let retrieved = store.get(&id).unwrap();
117        assert_eq!(retrieved.id.as_str(), "test-1");
118        assert!((retrieved.outcome - 0.8).abs() < 0.001);
119
120        // Contains
121        assert!(store.contains(&id));
122        assert!(!store.contains(&"nonexistent".into()));
123
124        // Length
125        assert_eq!(store.len(), 1);
126
127        // Update stats
128        store.update_stats(&id, 0.9).unwrap();
129        store.update_stats(&id, 0.7).unwrap();
130
131        let updated = store.get(&id).unwrap();
132        assert_eq!(updated.stats.count(), 2);
133
134        // Remove
135        assert!(store.remove(&id).unwrap());
136        assert!(!store.contains(&id));
137        assert_eq!(store.len(), 0);
138    }
139
140    pub fn test_batch_operations<S: RecordStore>(store: &mut S) {
141        use crate::types::RecordStatus;
142
143        let records: Vec<MemoryRecord> = (0..10)
144            .map(|i| MemoryRecord {
145                id: format!("batch-{i}").into(),
146                embedding: vec![i as f32; 3],
147                context: format!("context {i}"),
148                outcome: i as f64 / 10.0,
149                metadata: Default::default(),
150                created_at: i as u64,
151                status: RecordStatus::Active,
152                stats: OutcomeStats::new(1),
153            })
154            .collect();
155
156        let ids = store.insert_batch(records).unwrap();
157        assert_eq!(ids.len(), 10);
158        assert_eq!(store.len(), 10);
159
160        // Get batch
161        let retrieved = store.get_batch(&ids[..3]);
162        assert_eq!(retrieved.len(), 3);
163        assert!(retrieved.iter().all(|r| r.is_some()));
164
165        // Get IDs
166        let all_ids = store.ids();
167        assert_eq!(all_ids.len(), 10);
168
169        // Clear
170        store.clear();
171        assert!(store.is_empty());
172    }
173
174    pub fn test_duplicate_insert<S: RecordStore>(store: &mut S) {
175        use crate::types::RecordStatus;
176
177        let record = MemoryRecord {
178            id: "dup-test".into(),
179            embedding: vec![1.0],
180            context: "test".into(),
181            outcome: 0.5,
182            metadata: Default::default(),
183            created_at: 0,
184            status: RecordStatus::Active,
185            stats: OutcomeStats::new(1),
186        };
187
188        store.insert(record.clone()).unwrap();
189        let result = store.insert(record);
190        assert!(result.is_err());
191    }
192}
rag_plusplus_core/store/traits.rs

rag_plusplus_core/store/
traits.rs
