1use std::sync::Arc;
4use std::time::{SystemTime, UNIX_EPOCH};
5use uuid::Uuid;
6
7use crate::config::{Config, DEFAULT_RECENCY_HALF_LIFE_SECS};
8use crate::entry::ContextEntry;
9use crate::error::Error;
10use crate::lexicon::LexiconScorer;
11use crate::traits::{ContextStorage, Result, Searcher};
12#[cfg(feature = "semantic")]
14use crate::semantic::Embedder as _;
15
16const DEFAULT_SEARCH_LIMIT: usize = 50;
18
19pub const MATCH_ALL_QUERY: &str = "*";
24
25#[derive(Debug, Default, Clone)]
27pub struct SaveOptions {
28 pub session_id: Option<String>,
30 pub scope: Option<String>,
32 pub metadata: Option<serde_json::Value>,
39}
40
41#[must_use]
43pub fn estimate_tokens(text: &str) -> usize {
44 text.len().div_ceil(4)
45}
46
47fn recency_decay(age_seconds: f64, half_life: f64) -> f64 {
52 0.5_f64.powf(age_seconds / half_life)
53}
54
55pub struct ContextEngine {
57 storage: Box<dyn ContextStorage>,
58 searcher: Box<dyn Searcher>,
59 config: Config,
60 scorer: Option<Arc<dyn LexiconScorer>>,
61 #[cfg(feature = "semantic")]
62 embedder: Option<Arc<crate::semantic::FasEmbedder>>,
63}
64
65impl ContextEngine {
66 #[must_use]
72 pub fn new(
73 storage: Box<dyn ContextStorage>,
74 searcher: Box<dyn Searcher>,
75 mut config: Config,
76 ) -> Self {
77 if !config.recency_half_life_secs.is_finite() || config.recency_half_life_secs <= 0.0 {
78 config.recency_half_life_secs = DEFAULT_RECENCY_HALF_LIFE_SECS;
79 }
80
81 Self {
82 storage,
83 searcher,
84 config,
85 scorer: None,
86 #[cfg(feature = "semantic")]
87 embedder: None,
88 }
89 }
90
91 #[must_use]
99 pub fn with_scorer(mut self, scorer: Arc<dyn LexiconScorer>) -> Self {
100 self.scorer = Some(scorer);
101 self
102 }
103
104 #[cfg(feature = "semantic")]
110 #[must_use]
111 pub(crate) fn with_embedder(mut self, embedder: Arc<crate::semantic::FasEmbedder>) -> Self {
112 self.embedder = Some(embedder);
113 self
114 }
115
116 pub async fn assemble(
130 &self,
131 query: &str,
132 scope: Option<&str>,
133 token_budget: usize,
134 ) -> Result<Vec<ContextEntry>> {
135 let limit = DEFAULT_SEARCH_LIMIT;
136
137 tracing::debug!(query = %query, ?scope, token_budget, "assemble: bm25 search");
138 let bm25_candidates = self.searcher.search(query, scope, limit).await?;
139 tracing::debug!(count = %bm25_candidates.len(), "assemble: bm25 complete");
140
141 let semantic_candidates = self.run_semantic_search(query, scope, limit).await;
144 tracing::debug!(count = %semantic_candidates.len(), "assemble: semantic complete");
145
146 if bm25_candidates.is_empty() && semantic_candidates.is_empty() {
147 tracing::debug!("assemble: both searches empty");
148 return Ok(Vec::new());
149 }
150
151 let worst_rank = limit + 1;
157 let k = 60.0_f64;
158 let now = current_timestamp();
159 let half_life = self.config.recency_half_life_secs;
160
161 let mut entry_map: std::collections::HashMap<String, ContextEntry> =
163 std::collections::HashMap::new();
164 for se in &semantic_candidates {
165 entry_map
166 .entry(se.entry.id.clone())
167 .or_insert_with(|| se.entry.clone());
168 }
169 for se in &bm25_candidates {
170 entry_map.insert(se.entry.id.clone(), se.entry.clone());
171 }
172
173 let bm25_rank: std::collections::HashMap<&str, usize> = bm25_candidates
174 .iter()
175 .enumerate()
176 .map(|(i, se)| (se.entry.id.as_str(), i + 1))
177 .collect();
178 let semantic_rank: std::collections::HashMap<&str, usize> = semantic_candidates
179 .iter()
180 .enumerate()
181 .map(|(i, se)| (se.entry.id.as_str(), i + 1))
182 .collect();
183
184 let mut weighted: Vec<(f64, ContextEntry)> = entry_map
185 .into_iter()
186 .map(|(id, entry)| {
187 let br = *bm25_rank.get(id.as_str()).unwrap_or(&worst_rank);
188 let sr = *semantic_rank.get(id.as_str()).unwrap_or(&worst_rank);
189 #[allow(
190 clippy::cast_precision_loss,
191 reason = "ranks are small integers (max ~51); lossless"
192 )]
193 let rrf = 1.0 / (k + br as f64) + 1.0 / (k + sr as f64);
194
195 #[allow(
196 clippy::cast_precision_loss,
197 reason = "Unix timestamps fit losslessly in f64 for millions of years"
198 )]
199 let age = (now - entry.timestamp).max(0) as f64;
200 let decay = recency_decay(age, half_life);
201 let boost = self
202 .scorer
203 .as_ref()
204 .map_or(0.0_f32, |s| s.score(&entry, query));
205 let final_score = rrf * decay * (1.0 + f64::from(boost).clamp(-1.0, 2.0));
206 (final_score, entry)
207 })
208 .collect();
209
210 weighted.sort_by(|a, b| b.0.total_cmp(&a.0));
211 tracing::debug!(fused = %weighted.len(), "assemble: rrf fusion complete");
212
213 let mut result = Vec::new();
215 let mut tokens_used: usize = 0;
216 for (_score, entry) in weighted {
217 let entry_tokens = entry
218 .token_count
219 .unwrap_or_else(|| estimate_tokens(&entry.content));
220 if tokens_used.saturating_add(entry_tokens) > token_budget {
221 continue;
222 }
223 tokens_used += entry_tokens;
224 result.push(entry);
225 }
226
227 tracing::debug!(entries = %result.len(), tokens_used, "assemble: complete");
228 Ok(result)
229 }
230
231 #[allow(
233 clippy::unused_async,
234 reason = "async only active under semantic feature"
235 )]
236 async fn run_semantic_search(
237 &self,
238 query: &str,
239 scope: Option<&str>,
240 limit: usize,
241 ) -> Vec<crate::entry::ScoredEntry> {
242 #[cfg(feature = "semantic")]
243 if let Some(ref embedder) = self.embedder {
244 let emb = embedder.clone();
245 let query_owned = query.to_owned();
246 match tokio::task::spawn_blocking(move || emb.embed(&query_owned)).await {
247 Ok(Ok(embedding)) => {
248 tracing::debug!(dims = %embedding.len(), "query embedding ready");
249 match self
250 .searcher
251 .search_semantic(&embedding, scope, limit)
252 .await
253 {
254 Ok(results) => return results,
255 Err(e) => tracing::warn!(error = %e, "semantic search failed"),
256 }
257 }
258 Ok(Err(e)) => tracing::warn!(error = %e, "query embedding failed"),
259 Err(e) => tracing::warn!(error = %e, "embed task panicked"),
260 }
261 }
262 #[cfg(not(feature = "semantic"))]
264 let _ = (query, scope, limit);
265 Vec::new()
266 }
267
268 pub async fn save_snapshot(
287 &self,
288 content: &str,
289 kind: &str,
290 options: &SaveOptions,
291 ) -> Result<String> {
292 if content.is_empty() {
293 return Err(Error::InvalidEntry("content must not be empty".into()));
294 }
295
296 let timestamp = current_timestamp();
297 let id = Uuid::now_v7().to_string();
298 let token_count = estimate_tokens(content);
299
300 let entry = ContextEntry {
301 id: id.clone(),
302 content: content.to_owned(),
303 timestamp,
304 kind: kind.to_owned(),
305 scope: options.scope.clone(),
306 session_id: options.session_id.clone(),
307 token_count: Some(token_count),
308 metadata: options.metadata.clone(),
309 };
310
311 self.storage.save(&entry).await?;
312 tracing::trace!(id = %id, kind = %kind, tokens = %token_count, "entry saved");
313
314 self.embed_and_store(&id, content).await;
316
317 Ok(id)
318 }
319
320 #[allow(
325 clippy::unused_async,
326 reason = "async only active under semantic feature"
327 )]
328 async fn embed_and_store(&self, id: &str, content: &str) {
329 #[cfg(feature = "semantic")]
330 if let Some(ref embedder) = self.embedder {
331 tracing::debug!(id = %id, "generating embedding");
332 let emb = embedder.clone();
333 let text = content.to_owned();
334 let id_owned = id.to_owned();
335 match tokio::task::spawn_blocking(move || emb.embed(&text)).await {
336 Ok(Ok(embedding)) => {
337 if let Err(e) = self.storage.save_embedding(&id_owned, &embedding).await {
338 tracing::warn!(
339 id = %id_owned,
340 error = %e,
341 "save_embedding failed; entry stored without semantic index"
342 );
343 } else {
344 tracing::trace!(id = %id_owned, dims = %embedding.len(), "embedding stored");
345 }
346 }
347 Ok(Err(e)) => tracing::warn!(
348 id = %id,
349 error = %e,
350 "embed failed; entry stored without semantic index"
351 ),
352 Err(e) => tracing::warn!(
353 id = %id,
354 error = %e,
355 "embed task panicked; entry stored without semantic index"
356 ),
357 }
358 }
359 #[cfg(not(feature = "semantic"))]
360 let _ = (id, content);
361 }
362
363 #[must_use]
369 pub fn storage(&self) -> &dyn ContextStorage {
370 self.storage.as_ref()
371 }
372
373 #[cfg(feature = "semantic")]
386 pub async fn backfill_embeddings(
387 &self,
388 batch_size: usize,
389 progress: impl Fn(usize, usize),
390 ) -> Result<usize> {
391 let embedder = if let Some(e) = &self.embedder {
392 e.clone()
393 } else {
394 tracing::debug!("backfill_embeddings: no embedder configured");
395 return Ok(0);
396 };
397
398 let all = self.storage.get_unembedded(usize::MAX).await?;
399 let total = all.len();
400 if total == 0 {
401 tracing::debug!("backfill_embeddings: all entries already embedded");
402 return Ok(0);
403 }
404
405 tracing::debug!(total = %total, batch_size = %batch_size, "backfill_embeddings: starting");
406 let batch = batch_size.max(1);
407 let mut n_embedded = 0usize;
408
409 for chunk in all.chunks(batch) {
410 let texts: Vec<String> = chunk.iter().map(|e| e.content.clone()).collect();
411 let emb = embedder.clone();
412 let embeddings = tokio::task::spawn_blocking(move || {
413 let refs: Vec<&str> = texts.iter().map(String::as_str).collect();
414 emb.embed_batch(&refs)
415 })
416 .await
417 .map_err(|e| Error::Migration(format!("backfill embed task panicked: {e}")))?
418 .map_err(|e| Error::Migration(format!("backfill embed failed: {e}")))?;
419
420 for (entry, embedding) in chunk.iter().zip(embeddings.iter()) {
421 if let Err(e) = self.storage.save_embedding(&entry.id, embedding).await {
422 tracing::warn!(id = %entry.id, error = %e, "backfill: save_embedding failed");
423 } else {
424 n_embedded += 1;
425 }
426 }
427
428 progress(n_embedded, total);
429 tracing::debug!(done = %n_embedded, total = %total, "backfill_embeddings: batch done");
430 }
431
432 Ok(n_embedded)
433 }
434}
435
436fn current_timestamp() -> i64 {
441 let secs = SystemTime::now()
442 .duration_since(UNIX_EPOCH)
443 .map_or(0, |d| d.as_secs());
444 i64::try_from(secs).unwrap_or(i64::MAX)
445}
446
447#[cfg(test)]
448mod tests {
449 use super::*;
450 use crate::config::{EvictionPolicy, DEFAULT_RECENCY_HALF_LIFE_SECS};
451 use crate::entry::ScoredEntry;
452 use async_trait::async_trait;
453 use std::path::PathBuf;
454 use std::sync::Mutex;
455
456 struct MockStorage {
461 entries: Mutex<Vec<ContextEntry>>,
462 }
463
464 impl MockStorage {
465 fn new() -> Self {
466 Self {
467 entries: Mutex::new(Vec::new()),
468 }
469 }
470 }
471
472 #[async_trait]
473 impl ContextStorage for MockStorage {
474 async fn save(&self, entry: &ContextEntry) -> Result<()> {
475 self.entries.lock().unwrap().push(entry.clone());
476 Ok(())
477 }
478
479 async fn get_top_k(&self, k: usize) -> Result<Vec<ContextEntry>> {
480 let guard = self.entries.lock().unwrap();
481 let mut sorted = guard.clone();
482 sorted.sort_by_key(|e| std::cmp::Reverse(e.timestamp));
483 sorted.truncate(k);
484 Ok(sorted)
485 }
486
487 async fn get_all(&self) -> Result<Vec<ContextEntry>> {
488 Ok(self.entries.lock().unwrap().clone())
489 }
490
491 async fn delete(&self, id: &str) -> Result<bool> {
492 let mut guard = self.entries.lock().unwrap();
493 let before = guard.len();
494 guard.retain(|e| e.id != id);
495 Ok(guard.len() < before)
496 }
497
498 async fn clear(&self) -> Result<usize> {
499 let mut guard = self.entries.lock().unwrap();
500 let n = guard.len();
501 guard.clear();
502 Ok(n)
503 }
504
505 async fn clear_scope(&self, scope: &str) -> Result<usize> {
506 let mut guard = self.entries.lock().unwrap();
507 let before = guard.len();
508 guard.retain(|e| e.scope.as_deref() != Some(scope));
509 Ok(before - guard.len())
510 }
511
512 async fn count(&self) -> Result<usize> {
513 Ok(self.entries.lock().unwrap().len())
514 }
515 }
516
517 struct MockSearcher {
518 results: Mutex<Vec<ScoredEntry>>,
519 }
520
521 impl MockSearcher {
522 fn new(results: Vec<ScoredEntry>) -> Self {
523 Self {
524 results: Mutex::new(results),
525 }
526 }
527
528 fn empty() -> Self {
529 Self::new(Vec::new())
530 }
531 }
532
533 #[async_trait]
534 impl Searcher for MockSearcher {
535 async fn search(
536 &self,
537 _query: &str,
538 _scope: Option<&str>,
539 limit: usize,
540 ) -> Result<Vec<ScoredEntry>> {
541 let guard = self.results.lock().unwrap();
542 Ok(guard.iter().take(limit).cloned().collect())
543 }
544 }
545
546 fn default_config(max_entries: usize) -> Config {
547 Config {
548 max_entries,
549 token_budget: 8192,
550 db_path: PathBuf::from(":memory:"),
551 eviction_policy: EvictionPolicy::Lru,
552 recency_half_life_secs: DEFAULT_RECENCY_HALF_LIFE_SECS,
553 ..Config::default()
554 }
555 }
556
557 fn make_entry(id: &str, content: &str, timestamp: i64) -> ContextEntry {
558 ContextEntry {
559 id: id.into(),
560 content: content.into(),
561 timestamp,
562 kind: crate::entry::kind::MANUAL.to_owned(),
563 scope: None,
564 session_id: None,
565 token_count: Some(estimate_tokens(content)),
566 metadata: None,
567 }
568 }
569
570 fn make_scored(id: &str, content: &str, timestamp: i64, score: f64) -> ScoredEntry {
571 ScoredEntry {
572 entry: make_entry(id, content, timestamp),
573 score,
574 }
575 }
576
577 #[test]
582 fn test_estimate_tokens() {
583 assert_eq!(estimate_tokens(""), 0);
584 assert_eq!(estimate_tokens("a"), 1); assert_eq!(estimate_tokens("ab"), 1); assert_eq!(estimate_tokens("abc"), 1); assert_eq!(estimate_tokens("abcd"), 1); assert_eq!(estimate_tokens("abcde"), 2); assert_eq!(estimate_tokens("abcdefgh"), 2); assert_eq!(estimate_tokens("hello world, this is a test"), 7);
591 }
592
593 #[test]
594 fn test_engine_send_sync() {
595 fn assert_send_sync<T: Send + Sync>() {}
596 assert_send_sync::<ContextEngine>();
597 }
598
599 #[tokio::test]
600 async fn test_assemble_fits_budget() {
601 let now = current_timestamp();
602 let results = vec![
603 make_scored("a", "short", now, 1.0),
604 make_scored("b", "medium length text here", now, 0.9),
605 make_scored("c", "another entry with more content inside", now, 0.8),
606 ];
607
608 let engine = ContextEngine::new(
609 Box::new(MockStorage::new()),
610 Box::new(MockSearcher::new(results)),
611 default_config(100),
612 );
613
614 let assembled = engine.assemble("test", None, 2).await.unwrap();
616 assert_eq!(assembled.len(), 1);
617 assert_eq!(assembled[0].id, "a");
618
619 let assembled = engine.assemble("test", None, 1000).await.unwrap();
621 assert_eq!(assembled.len(), 3);
622 }
623
624 #[tokio::test]
625 async fn test_assemble_skips_oversized_entries() {
626 let now = current_timestamp();
628 let big_content = "x".repeat(4000); let results = vec![
630 make_scored("big", &big_content, now, 1.0),
631 make_scored("small", "fits", now, 0.9),
632 ];
633
634 let engine = ContextEngine::new(
635 Box::new(MockStorage::new()),
636 Box::new(MockSearcher::new(results)),
637 default_config(100),
638 );
639
640 let assembled = engine.assemble("test", None, 5).await.unwrap();
642 assert_eq!(assembled.len(), 1);
643 assert_eq!(assembled[0].id, "small");
644 }
645
646 #[tokio::test]
647 async fn test_assemble_empty_results() {
648 let engine = ContextEngine::new(
649 Box::new(MockStorage::new()),
650 Box::new(MockSearcher::empty()),
651 default_config(100),
652 );
653 let assembled = engine.assemble("anything", None, 1000).await.unwrap();
654 assert!(assembled.is_empty());
655 }
656
657 #[tokio::test]
658 async fn test_recency_weighting() {
659 let now = current_timestamp();
660 let results = vec![
662 make_scored("old", "old entry", now - 86_400, 1.0), make_scored("new", "new entry", now, 1.0), ];
665
666 let engine = ContextEngine::new(
667 Box::new(MockStorage::new()),
668 Box::new(MockSearcher::new(results)),
669 default_config(100),
670 );
671
672 let assembled = engine.assemble("test", None, 1000).await.unwrap();
673 assert_eq!(assembled.len(), 2);
674 assert_eq!(assembled[0].id, "new");
676 assert_eq!(assembled[1].id, "old");
677 }
678
679 #[tokio::test]
680 async fn test_save_snapshot_creates_entry() {
681 let engine = ContextEngine::new(
682 Box::new(MockStorage::new()),
683 Box::new(MockSearcher::empty()),
684 default_config(100),
685 );
686
687 let id = engine
688 .save_snapshot(
689 "hello world",
690 crate::entry::kind::MANUAL,
691 &SaveOptions::default(),
692 )
693 .await
694 .unwrap();
695 assert!(!id.is_empty());
696
697 let all = engine.storage.get_all().await.unwrap();
699 assert_eq!(all.len(), 1);
700 assert_eq!(all[0].id, id);
701 assert_eq!(all[0].content, "hello world");
702 assert_eq!(all[0].token_count, Some(estimate_tokens("hello world")));
703 assert_eq!(all[0].kind, crate::entry::kind::MANUAL);
704 }
705
706 #[tokio::test]
707 async fn test_save_snapshot_populates_scope_and_metadata() {
708 let engine = ContextEngine::new(
709 Box::new(MockStorage::new()),
710 Box::new(MockSearcher::empty()),
711 default_config(100),
712 );
713
714 let metadata = serde_json::json!({"source": "test"});
715 let options = SaveOptions {
716 session_id: Some("sess-1".to_owned()),
717 scope: Some("project:homelab-rs".to_owned()),
718 metadata: Some(metadata.clone()),
719 };
720
721 let id = engine
722 .save_snapshot("hello scoped", crate::entry::kind::SNAPSHOT, &options)
723 .await
724 .unwrap();
725
726 let all = engine.storage.get_all().await.unwrap();
727 assert_eq!(all.len(), 1);
728 assert_eq!(all[0].id, id);
729 assert_eq!(all[0].kind, crate::entry::kind::SNAPSHOT);
730 assert_eq!(all[0].scope.as_deref(), Some("project:homelab-rs"));
731 assert_eq!(all[0].session_id.as_deref(), Some("sess-1"));
732 assert_eq!(all[0].metadata, Some(metadata));
733 }
734
735 #[test]
736 fn test_recency_decay_values() {
737 let half_life = crate::config::DEFAULT_RECENCY_HALF_LIFE_HOURS * 3600.0;
738
739 let decay_0 = recency_decay(0.0, half_life);
741 assert!((decay_0 - 1.0).abs() < f64::EPSILON);
742
743 let decay_half = recency_decay(half_life, half_life);
745 assert!((decay_half - 0.5).abs() < 1e-10);
746
747 let decay_double = recency_decay(2.0 * half_life, half_life);
749 assert!((decay_double - 0.25).abs() < 1e-10);
750 }
751
752 #[tokio::test]
753 async fn test_save_snapshot_ids_are_unique() {
754 let engine = ContextEngine::new(
755 Box::new(MockStorage::new()),
756 Box::new(MockSearcher::empty()),
757 default_config(100),
758 );
759
760 let id1 = engine
761 .save_snapshot(
762 "identical content",
763 crate::entry::kind::SNAPSHOT,
764 &SaveOptions::default(),
765 )
766 .await
767 .unwrap();
768 let id2 = engine
769 .save_snapshot(
770 "identical content",
771 crate::entry::kind::SNAPSHOT,
772 &SaveOptions::default(),
773 )
774 .await
775 .unwrap();
776 assert_ne!(
777 id1, id2,
778 "Two saves of identical content must produce distinct IDs"
779 );
780 }
781
782 #[tokio::test]
783 async fn test_save_empty_content_rejected() {
784 let engine = ContextEngine::new(
785 Box::new(MockStorage::new()),
786 Box::new(MockSearcher::empty()),
787 default_config(100),
788 );
789
790 let result = engine
791 .save_snapshot("", crate::entry::kind::MANUAL, &SaveOptions::default())
792 .await;
793 assert!(result.is_err());
794 let err = result.unwrap_err();
795 assert!(err.to_string().contains("content must not be empty"));
796 }
797
798 #[test]
799 fn test_invalid_half_life_clamped_to_default() {
800 let invalid_values: Vec<f64> = vec![
801 0.0,
802 -1.0,
803 -100.0,
804 f64::NAN,
805 f64::INFINITY,
806 f64::NEG_INFINITY,
807 ];
808
809 for value in invalid_values {
810 let mut config = default_config(100);
811 config.recency_half_life_secs = value;
812
813 let engine = ContextEngine::new(
814 Box::new(MockStorage::new()),
815 Box::new(MockSearcher::empty()),
816 config,
817 );
818
819 assert!(
820 (engine.config.recency_half_life_secs - DEFAULT_RECENCY_HALF_LIFE_SECS).abs()
821 < f64::EPSILON,
822 "half_life {value} should have been clamped to default",
823 );
824 }
825 }
826
827 #[tokio::test]
828 async fn lexicon_scorer_boosts_important_entry_over_neutral() {
829 use crate::lexicon::DefaultEnglishScorer;
830
831 let now = current_timestamp();
832 let results = vec![
834 make_scored("neutral", "something unrelated here", now, 1.0),
835 make_scored("important", "confirmed, that is correct", now, 1.0),
836 ];
837
838 let scorer: Arc<dyn LexiconScorer> = Arc::new(DefaultEnglishScorer::default());
839 let engine = ContextEngine::new(
840 Box::new(MockStorage::new()),
841 Box::new(MockSearcher::new(results)),
842 default_config(100),
843 )
844 .with_scorer(scorer);
845
846 let assembled = engine.assemble("test", None, 1000).await.unwrap();
847 assert_eq!(assembled.len(), 2);
848 assert_eq!(
849 assembled[0].id, "important",
850 "lexicon-boosted entry should rank first"
851 );
852 }
853
854 #[test]
855 fn test_valid_half_life_preserved() {
856 let mut config = default_config(100);
857 config.recency_half_life_secs = 3600.0; let engine = ContextEngine::new(
860 Box::new(MockStorage::new()),
861 Box::new(MockSearcher::empty()),
862 config,
863 );
864
865 assert!((engine.config.recency_half_life_secs - 3600.0).abs() < f64::EPSILON);
866 }
867}