1use crate::vector_search;
4
5#[derive(Debug, Clone)]
7pub struct MemoryCache {
8 pub chunks: Vec<String>,
9 pub embeddings: Vec<Vec<f32>>,
10 pub source_channels: Vec<String>,
11 pub timestamps: Vec<f64>,
12 pub session_ids: Vec<String>,
13 pub tags: Vec<String>,
14 pub tombstones: Vec<u8>,
15 pub embedding_dim: usize,
16 pub norms: Vec<f32>,
18 pub activation_weights: Vec<f32>,
20}
21
22impl MemoryCache {
23 pub fn new(embedding_dim: usize) -> Self {
24 Self {
25 chunks: Vec::new(),
26 embeddings: Vec::new(),
27 source_channels: Vec::new(),
28 timestamps: Vec::new(),
29 session_ids: Vec::new(),
30 tags: Vec::new(),
31 tombstones: Vec::new(),
32 embedding_dim,
33 norms: Vec::new(),
34 activation_weights: Vec::new(),
35 }
36 }
37
38 pub fn len(&self) -> usize {
40 self.chunks.len()
41 }
42
43 pub fn is_empty(&self) -> bool {
44 self.chunks.is_empty()
45 }
46
47 pub fn count_active(&self) -> usize {
49 self.tombstones.iter().filter(|&&t| t == 0).count()
50 }
51
52 pub fn push(
54 &mut self,
55 chunk: String,
56 embedding: Vec<f32>,
57 source_channel: String,
58 timestamp: f64,
59 session_id: String,
60 tags: String,
61 ) -> usize {
62 let idx = self.chunks.len();
63 let norm = vector_search::compute_norm(&embedding);
64 self.chunks.push(chunk);
65 self.embeddings.push(embedding);
66 self.source_channels.push(source_channel);
67 self.timestamps.push(timestamp);
68 self.session_ids.push(session_id);
69 self.tags.push(tags);
70 self.tombstones.push(0);
71 self.norms.push(norm);
72 self.activation_weights.push(1.0);
73 idx
74 }
75
76 pub fn mark_deleted(&mut self, id: usize) -> bool {
78 if id < self.tombstones.len() && self.tombstones[id] == 0 {
79 self.tombstones[id] = 1;
80 true
81 } else {
82 false
83 }
84 }
85
86 pub fn tombstone_fraction(&self) -> f32 {
88 if self.chunks.is_empty() {
89 return 0.0;
90 }
91 let tombstoned = self.tombstones.iter().filter(|&&t| t == 1).count();
92 tombstoned as f32 / self.chunks.len() as f32
93 }
94
95 pub fn compact(&mut self) -> (usize, Vec<Option<usize>>) {
99 let old_len = self.chunks.len();
100 let mut index_map = vec![None; old_len];
101 let mut new_idx = 0usize;
102
103 let mut new_chunks = Vec::new();
104 let mut new_embeddings = Vec::new();
105 let mut new_source_channels = Vec::new();
106 let mut new_timestamps = Vec::new();
107 let mut new_session_ids = Vec::new();
108 let mut new_tags = Vec::new();
109 let mut new_tombstones = Vec::new();
110 let mut new_norms = Vec::new();
111 let mut new_activation_weights = Vec::new();
112
113 for (i, slot) in index_map.iter_mut().enumerate() {
114 if self.tombstones[i] == 0 {
115 *slot = Some(new_idx);
116 new_idx += 1;
117 let norm = vector_search::compute_norm(&self.embeddings[i]);
118 new_chunks.push(self.chunks[i].clone());
119 new_embeddings.push(self.embeddings[i].clone());
120 new_source_channels.push(self.source_channels[i].clone());
121 new_timestamps.push(self.timestamps[i]);
122 new_session_ids.push(self.session_ids[i].clone());
123 new_tags.push(self.tags[i].clone());
124 new_tombstones.push(0u8);
125 new_norms.push(norm);
126 new_activation_weights.push(self.activation_weights[i]);
127 }
128 }
129
130 let removed = old_len - new_chunks.len();
131 self.chunks = new_chunks;
132 self.embeddings = new_embeddings;
133 self.source_channels = new_source_channels;
134 self.timestamps = new_timestamps;
135 self.session_ids = new_session_ids;
136 self.tags = new_tags;
137 self.tombstones = new_tombstones;
138 self.norms = new_norms;
139 self.activation_weights = new_activation_weights;
140
141 (removed, index_map)
142 }
143
144 pub fn flat_embeddings(&self) -> Vec<f32> {
146 let mut flat = Vec::with_capacity(self.embeddings.len() * self.embedding_dim);
147 for emb in &self.embeddings {
148 flat.extend_from_slice(emb);
149 }
150 flat
151 }
152}