oxios_memory/memory/manager/
mod.rs1mod ops;
12mod store;
13
14use std::collections::HashMap;
15use std::sync::Arc;
16
17use anyhow::Result;
18use parking_lot::RwLock;
19
20use crate::memory::embedding::{EmbeddingProvider, EmbeddingVector, TfIdfEmbeddingProvider};
21use crate::memory::hnsw_memory_index::HnswMemoryIndex;
22use crate::memory::sona::SonaEngine;
23use crate::memory::storage::{MemoryGit, MemoryStorage};
24use crate::memory::types::{MemoryEntry, MemoryType};
25
26use super::{CurationCandidate, CurationReport, MemoryBudget};
27
28pub struct MemoryManager {
38 pub(crate) storage: Arc<dyn MemoryStorage>,
40 pub(crate) max_recall: usize,
42 pub(crate) vector_index: RwLock<HashMap<String, EmbeddingVector>>,
44 pub(crate) embedding: Arc<dyn EmbeddingProvider>,
46 pub(crate) git: Option<Arc<dyn MemoryGit>>,
48 pub(crate) hnsw_index: RwLock<Option<Arc<HnswMemoryIndex>>>,
50 pub(crate) sona_engine: Option<Arc<SonaEngine>>,
52 #[cfg(feature = "sqlite-memory")]
54 pub(crate) sqlite_store: Option<Arc<crate::memory::sqlite::SqliteMemoryStore>>,
55}
56
57impl std::fmt::Debug for MemoryManager {
58 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
59 f.debug_struct("MemoryManager")
60 .field("max_recall", &self.max_recall)
61 .field("index_size", &self.vector_index.read().len())
62 .field("sona_enabled", &self.sona_engine.is_some())
63 .finish()
64 }
65}
66
67impl MemoryManager {
68 pub fn new(storage: Arc<dyn MemoryStorage>) -> Self {
70 Self {
71 storage,
72 max_recall: 10,
73 vector_index: RwLock::new(HashMap::new()),
74 embedding: Arc::new(TfIdfEmbeddingProvider),
75 git: None,
76 hnsw_index: RwLock::new(None),
77 sona_engine: None,
78 #[cfg(feature = "sqlite-memory")]
79 sqlite_store: None,
80 }
81 }
82
83 pub fn set_git_layer(&mut self, gl: Arc<dyn MemoryGit>) {
85 self.git = Some(gl);
86 }
87
88 #[cfg(feature = "sqlite-memory")]
90 pub fn set_sqlite_store(&mut self, store: Arc<crate::memory::sqlite::SqliteMemoryStore>) {
91 self.sqlite_store = Some(store);
92 }
93
94 #[cfg(feature = "sqlite-memory")]
96 pub fn sqlite_store(&self) -> &Option<Arc<crate::memory::sqlite::SqliteMemoryStore>> {
97 &self.sqlite_store
98 }
99
100 pub fn set_sona_engine(&mut self, engine: Arc<SonaEngine>) {
102 self.sona_engine = Some(engine);
103 }
104
105 pub fn sona_engine(&self) -> Option<&Arc<SonaEngine>> {
107 self.sona_engine.as_ref()
108 }
109
110 pub fn set_hnsw_index(&self, index: Arc<HnswMemoryIndex>) {
112 *self.hnsw_index.write() = Some(index);
113 }
114
115 pub fn with_max_recall(mut self, n: usize) -> Self {
117 self.max_recall = n;
118 self
119 }
120
121 pub fn set_max_recall(&mut self, n: usize) {
123 self.max_recall = n;
124 }
125
126 pub fn vector_index_size(&self) -> usize {
128 self.vector_index.read().len()
129 }
130
131 pub(crate) async fn git_commit(&self, rel_path: &str, message: &str) {
133 if let Some(ref gl) = self.git {
134 if gl.is_enabled() {
135 let _ = gl.commit_file(rel_path, message).await;
136 }
137 }
138 }
139
140 pub fn effective_importance(entry: &MemoryEntry) -> f32 {
142 let access_boost = (1.0_f32 + entry.access_count as f32).ln();
143 entry.importance * (1.0 + access_boost)
144 }
145
146 pub async fn curate(&self, budget: &MemoryBudget) -> Result<CurationReport> {
148 let mut report = CurationReport::default();
149
150 for mt in &[
151 MemoryType::Conversation,
152 MemoryType::Session,
153 MemoryType::Fact,
154 MemoryType::Episode,
155 MemoryType::Knowledge,
156 ] {
157 let entries = self.list(*mt, budget.max_per_type * 2).await?;
158 if entries.len() <= budget.max_per_type {
159 continue;
160 }
161
162 let total_count = entries.len();
163 let mut scored: Vec<_> = entries
164 .into_iter()
165 .map(|e| (e.id.clone(), e.memory_type, Self::effective_importance(&e)))
166 .collect();
167 scored.sort_by(|a, b| a.2.partial_cmp(&b.2).unwrap_or(std::cmp::Ordering::Equal));
168
169 let to_remove = scored.len() - budget.max_per_type;
170 for (id, memory_type, score) in scored.into_iter().take(to_remove) {
171 report.candidates_for_removal.push(CurationCandidate {
172 id,
173 memory_type,
174 effective_importance: score,
175 });
176 }
177 report.total_before += total_count;
178 }
179
180 for candidate in &report.candidates_for_removal {
181 if self
182 .forget(&candidate.id, candidate.memory_type)
183 .await
184 .is_ok()
185 {
186 report.removed += 1;
187 }
188 }
189
190 report.total_after = report.total_before - report.removed;
191 Ok(report)
192 }
193
194 pub fn spawn_curation_task(self: &Arc<Self>, budget: MemoryBudget) {
196 let mgr = Arc::clone(self);
197 tokio::spawn(async move {
198 match mgr.curate(&budget).await {
199 Ok(report) => {
200 if report.removed > 0 {
201 tracing::info!(
202 removed = report.removed,
203 candidates = report.candidates_for_removal.len(),
204 "Memory curation complete"
205 );
206 }
207 }
208 Err(e) => {
209 tracing::warn!(error = %e, "Memory curation failed");
210 }
211 }
212 });
213 }
214}