1mod algorithms;
5mod corrections;
6mod cross_session;
7mod graph;
8pub(crate) mod importance;
9pub mod persona;
10mod recall;
11mod summarization;
12pub mod trajectory;
13pub mod tree_consolidation;
14pub(crate) mod write_buffer;
15
16#[cfg(test)]
17mod tests;
18
19use std::sync::Arc;
20use std::sync::atomic::AtomicU64;
21
22use zeph_llm::any::AnyProvider;
23
24use crate::admission::AdmissionControl;
25use crate::embedding_store::EmbeddingStore;
26use crate::error::MemoryError;
27use crate::store::SqliteStore;
28use crate::token_counter::TokenCounter;
29
30pub(crate) const SESSION_SUMMARIES_COLLECTION: &str = "zeph_session_summaries";
31pub(crate) const KEY_FACTS_COLLECTION: &str = "zeph_key_facts";
32pub(crate) const CORRECTIONS_COLLECTION: &str = "zeph_corrections";
33
34pub use algorithms::{apply_mmr, apply_temporal_decay};
35pub use cross_session::SessionSummaryResult;
36pub use graph::{
37 ExtractionResult, ExtractionStats, GraphExtractionConfig, LinkingStats, NoteLinkingConfig,
38 PostExtractValidator, extract_and_store, link_memory_notes,
39};
40pub use persona::{
41 PersonaExtractionConfig, contains_self_referential_language, extract_persona_facts,
42};
43pub use recall::{EmbedContext, RecalledMessage};
44pub use summarization::{StructuredSummary, Summary, build_summarization_prompt};
45pub use trajectory::{TrajectoryEntry, TrajectoryExtractionConfig, extract_trajectory_entries};
46pub use tree_consolidation::{
47 TreeConsolidationConfig, TreeConsolidationResult, run_tree_consolidation_sweep,
48 start_tree_consolidation_loop,
49};
50pub use write_buffer::{BufferedWrite, WriteBuffer};
51
52pub struct SemanticMemory {
53 pub(crate) sqlite: SqliteStore,
54 pub(crate) qdrant: Option<Arc<EmbeddingStore>>,
55 pub(crate) provider: AnyProvider,
56 pub(crate) embed_provider: Option<AnyProvider>,
62 pub(crate) embedding_model: String,
63 pub(crate) vector_weight: f64,
64 pub(crate) keyword_weight: f64,
65 pub(crate) temporal_decay_enabled: bool,
66 pub(crate) temporal_decay_half_life_days: u32,
67 pub(crate) mmr_enabled: bool,
68 pub(crate) mmr_lambda: f32,
69 pub(crate) importance_enabled: bool,
70 pub(crate) importance_weight: f64,
71 pub(crate) tier_boost_semantic: f64,
74 pub token_counter: Arc<TokenCounter>,
75 pub graph_store: Option<Arc<crate::graph::GraphStore>>,
76 pub(crate) community_detection_failures: Arc<AtomicU64>,
77 pub(crate) graph_extraction_count: Arc<AtomicU64>,
78 pub(crate) graph_extraction_failures: Arc<AtomicU64>,
79 pub(crate) admission_control: Option<Arc<AdmissionControl>>,
81 pub(crate) key_facts_dedup_threshold: f32,
85}
86
87impl SemanticMemory {
88 pub async fn new(
99 sqlite_path: &str,
100 qdrant_url: &str,
101 provider: AnyProvider,
102 embedding_model: &str,
103 ) -> Result<Self, MemoryError> {
104 Self::with_weights(sqlite_path, qdrant_url, provider, embedding_model, 0.7, 0.3).await
105 }
106
107 pub async fn with_weights(
116 sqlite_path: &str,
117 qdrant_url: &str,
118 provider: AnyProvider,
119 embedding_model: &str,
120 vector_weight: f64,
121 keyword_weight: f64,
122 ) -> Result<Self, MemoryError> {
123 Self::with_weights_and_pool_size(
124 sqlite_path,
125 qdrant_url,
126 provider,
127 embedding_model,
128 vector_weight,
129 keyword_weight,
130 5,
131 )
132 .await
133 }
134
135 pub async fn with_weights_and_pool_size(
144 sqlite_path: &str,
145 qdrant_url: &str,
146 provider: AnyProvider,
147 embedding_model: &str,
148 vector_weight: f64,
149 keyword_weight: f64,
150 pool_size: u32,
151 ) -> Result<Self, MemoryError> {
152 let sqlite = SqliteStore::with_pool_size(sqlite_path, pool_size).await?;
153 let pool = sqlite.pool().clone();
154
155 let qdrant = match EmbeddingStore::new(qdrant_url, pool) {
156 Ok(store) => Some(Arc::new(store)),
157 Err(e) => {
158 tracing::warn!("Qdrant unavailable, semantic search disabled: {e:#}");
159 None
160 }
161 };
162
163 Ok(Self {
164 sqlite,
165 qdrant,
166 provider,
167 embed_provider: None,
168 embedding_model: embedding_model.into(),
169 vector_weight,
170 keyword_weight,
171 temporal_decay_enabled: false,
172 temporal_decay_half_life_days: 30,
173 mmr_enabled: false,
174 mmr_lambda: 0.7,
175 importance_enabled: false,
176 importance_weight: 0.15,
177 tier_boost_semantic: 1.3,
178 token_counter: Arc::new(TokenCounter::new()),
179 graph_store: None,
180 community_detection_failures: Arc::new(AtomicU64::new(0)),
181 graph_extraction_count: Arc::new(AtomicU64::new(0)),
182 graph_extraction_failures: Arc::new(AtomicU64::new(0)),
183 admission_control: None,
184 key_facts_dedup_threshold: 0.95,
185 })
186 }
187
188 pub async fn with_qdrant_ops(
197 sqlite_path: &str,
198 ops: crate::QdrantOps,
199 provider: AnyProvider,
200 embedding_model: &str,
201 vector_weight: f64,
202 keyword_weight: f64,
203 pool_size: u32,
204 ) -> Result<Self, MemoryError> {
205 let sqlite = SqliteStore::with_pool_size(sqlite_path, pool_size).await?;
206 let pool = sqlite.pool().clone();
207 let store = EmbeddingStore::with_store(Box::new(ops), pool);
208
209 Ok(Self {
210 sqlite,
211 qdrant: Some(Arc::new(store)),
212 provider,
213 embed_provider: None,
214 embedding_model: embedding_model.into(),
215 vector_weight,
216 keyword_weight,
217 temporal_decay_enabled: false,
218 temporal_decay_half_life_days: 30,
219 mmr_enabled: false,
220 mmr_lambda: 0.7,
221 importance_enabled: false,
222 importance_weight: 0.15,
223 tier_boost_semantic: 1.3,
224 token_counter: Arc::new(TokenCounter::new()),
225 graph_store: None,
226 community_detection_failures: Arc::new(AtomicU64::new(0)),
227 graph_extraction_count: Arc::new(AtomicU64::new(0)),
228 graph_extraction_failures: Arc::new(AtomicU64::new(0)),
229 admission_control: None,
230 key_facts_dedup_threshold: 0.95,
231 })
232 }
233
234 #[must_use]
239 pub fn with_graph_store(mut self, store: Arc<crate::graph::GraphStore>) -> Self {
240 self.graph_store = Some(store);
241 self
242 }
243
244 #[must_use]
246 pub fn community_detection_failures(&self) -> u64 {
247 use std::sync::atomic::Ordering;
248 self.community_detection_failures.load(Ordering::Relaxed)
249 }
250
251 #[must_use]
253 pub fn graph_extraction_count(&self) -> u64 {
254 use std::sync::atomic::Ordering;
255 self.graph_extraction_count.load(Ordering::Relaxed)
256 }
257
258 #[must_use]
260 pub fn graph_extraction_failures(&self) -> u64 {
261 use std::sync::atomic::Ordering;
262 self.graph_extraction_failures.load(Ordering::Relaxed)
263 }
264
265 #[must_use]
267 pub fn with_ranking_options(
268 mut self,
269 temporal_decay_enabled: bool,
270 temporal_decay_half_life_days: u32,
271 mmr_enabled: bool,
272 mmr_lambda: f32,
273 ) -> Self {
274 self.temporal_decay_enabled = temporal_decay_enabled;
275 self.temporal_decay_half_life_days = temporal_decay_half_life_days;
276 self.mmr_enabled = mmr_enabled;
277 self.mmr_lambda = mmr_lambda;
278 self
279 }
280
281 #[must_use]
283 pub fn with_importance_options(mut self, enabled: bool, weight: f64) -> Self {
284 self.importance_enabled = enabled;
285 self.importance_weight = weight;
286 self
287 }
288
289 #[must_use]
293 pub fn with_tier_boost(mut self, boost: f64) -> Self {
294 self.tier_boost_semantic = boost;
295 self
296 }
297
298 #[must_use]
303 pub fn with_admission_control(mut self, control: AdmissionControl) -> Self {
304 self.admission_control = Some(Arc::new(control));
305 self
306 }
307
308 #[must_use]
313 pub fn with_key_facts_dedup_threshold(mut self, threshold: f32) -> Self {
314 self.key_facts_dedup_threshold = threshold;
315 self
316 }
317
318 #[must_use]
324 pub fn with_embed_provider(mut self, embed_provider: AnyProvider) -> Self {
325 self.embed_provider = Some(embed_provider);
326 self
327 }
328
329 pub(crate) fn effective_embed_provider(&self) -> &AnyProvider {
333 self.embed_provider.as_ref().unwrap_or(&self.provider)
334 }
335
336 #[must_use]
340 pub fn from_parts(
341 sqlite: SqliteStore,
342 qdrant: Option<Arc<EmbeddingStore>>,
343 provider: AnyProvider,
344 embedding_model: impl Into<String>,
345 vector_weight: f64,
346 keyword_weight: f64,
347 token_counter: Arc<TokenCounter>,
348 ) -> Self {
349 Self {
350 sqlite,
351 qdrant,
352 provider,
353 embed_provider: None,
354 embedding_model: embedding_model.into(),
355 vector_weight,
356 keyword_weight,
357 temporal_decay_enabled: false,
358 temporal_decay_half_life_days: 30,
359 mmr_enabled: false,
360 mmr_lambda: 0.7,
361 importance_enabled: false,
362 importance_weight: 0.15,
363 tier_boost_semantic: 1.3,
364 token_counter,
365 graph_store: None,
366 community_detection_failures: Arc::new(AtomicU64::new(0)),
367 graph_extraction_count: Arc::new(AtomicU64::new(0)),
368 graph_extraction_failures: Arc::new(AtomicU64::new(0)),
369 admission_control: None,
370 key_facts_dedup_threshold: 0.95,
371 }
372 }
373
374 pub async fn with_sqlite_backend(
380 sqlite_path: &str,
381 provider: AnyProvider,
382 embedding_model: &str,
383 vector_weight: f64,
384 keyword_weight: f64,
385 ) -> Result<Self, MemoryError> {
386 Self::with_sqlite_backend_and_pool_size(
387 sqlite_path,
388 provider,
389 embedding_model,
390 vector_weight,
391 keyword_weight,
392 5,
393 )
394 .await
395 }
396
397 pub async fn with_sqlite_backend_and_pool_size(
403 sqlite_path: &str,
404 provider: AnyProvider,
405 embedding_model: &str,
406 vector_weight: f64,
407 keyword_weight: f64,
408 pool_size: u32,
409 ) -> Result<Self, MemoryError> {
410 let sqlite = SqliteStore::with_pool_size(sqlite_path, pool_size).await?;
411 let pool = sqlite.pool().clone();
412 let store = EmbeddingStore::new_sqlite(pool);
413
414 Ok(Self {
415 sqlite,
416 qdrant: Some(Arc::new(store)),
417 provider,
418 embed_provider: None,
419 embedding_model: embedding_model.into(),
420 vector_weight,
421 keyword_weight,
422 temporal_decay_enabled: false,
423 temporal_decay_half_life_days: 30,
424 mmr_enabled: false,
425 mmr_lambda: 0.7,
426 importance_enabled: false,
427 importance_weight: 0.15,
428 tier_boost_semantic: 1.3,
429 token_counter: Arc::new(TokenCounter::new()),
430 graph_store: None,
431 community_detection_failures: Arc::new(AtomicU64::new(0)),
432 graph_extraction_count: Arc::new(AtomicU64::new(0)),
433 graph_extraction_failures: Arc::new(AtomicU64::new(0)),
434 admission_control: None,
435 key_facts_dedup_threshold: 0.95,
436 })
437 }
438
439 #[must_use]
441 pub fn sqlite(&self) -> &SqliteStore {
442 &self.sqlite
443 }
444
445 pub async fn is_vector_store_connected(&self) -> bool {
450 match self.qdrant.as_ref() {
451 Some(store) => store.health_check().await,
452 None => false,
453 }
454 }
455
456 #[must_use]
458 pub fn has_vector_store(&self) -> bool {
459 self.qdrant.is_some()
460 }
461
462 #[must_use]
464 pub fn embedding_store(&self) -> Option<&Arc<EmbeddingStore>> {
465 self.qdrant.as_ref()
466 }
467
468 pub fn provider(&self) -> &AnyProvider {
470 &self.provider
471 }
472
473 pub async fn message_count(
479 &self,
480 conversation_id: crate::types::ConversationId,
481 ) -> Result<i64, MemoryError> {
482 self.sqlite.count_messages(conversation_id).await
483 }
484
485 pub async fn unsummarized_message_count(
491 &self,
492 conversation_id: crate::types::ConversationId,
493 ) -> Result<i64, MemoryError> {
494 let after_id = self
495 .sqlite
496 .latest_summary_last_message_id(conversation_id)
497 .await?
498 .unwrap_or(crate::types::MessageId(0));
499 self.sqlite
500 .count_messages_after(conversation_id, after_id)
501 .await
502 }
503}