1use lmdb::{Cursor, Database, Environment, Transaction, WriteFlags};
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use std::sync::{Arc, RwLock};
7use wm_core::{
8 CoreError, EpisodicCapturePolicy, EpisodicId, EpisodicKind, EpisodicRecord, MemoryTransition,
9 Result, ValidityState,
10};
11
12use crate::embedder::Embedder;
13use crate::enrichment::VocabularyEnrichment;
14use crate::episodic_keys::{AdaptiveAliases, key_index_terms_with_aliases};
15use crate::query_planner::QueryPlan;
16use crate::search::strip_stopwords;
17
18#[derive(Debug, Clone)]
20pub struct EpisodicSearchResult {
21 pub record: EpisodicRecord,
22 pub score: f32,
23 pub matched_terms: usize,
24}
25
26pub struct EpisodicStore<'a> {
28 env: &'a Environment,
29 db: Database,
30 term_db: Database,
31 term_cache: Arc<RwLock<HashMap<String, Vec<EpisodicId>>>>,
32 mutation_count: &'a std::sync::atomic::AtomicU64,
33 embedder: Option<Arc<dyn Embedder + Send + Sync>>,
34 aliases: Option<AdaptiveAliases>,
35 enrichment: Option<VocabularyEnrichment>,
36}
37
38impl<'a> EpisodicStore<'a> {
39 pub(crate) fn new(
40 env: &'a Environment,
41 db: Database,
42 term_db: Database,
43 term_cache: Arc<RwLock<HashMap<String, Vec<EpisodicId>>>>,
44 mutation_count: &'a std::sync::atomic::AtomicU64,
45 ) -> Self {
46 Self {
47 env,
48 db,
49 term_db,
50 term_cache,
51 mutation_count,
52 embedder: None,
53 aliases: None,
54 enrichment: None,
55 }
56 }
57
58 #[must_use]
60 pub fn with_adaptive_aliases(mut self, aliases: AdaptiveAliases) -> Self {
61 if !aliases.is_empty() {
62 self.aliases = Some(aliases);
63 }
64 self
65 }
66
67 #[must_use]
69 pub fn with_enrichment(mut self, enrichment: VocabularyEnrichment) -> Self {
70 if !enrichment.is_empty() {
71 self.enrichment = Some(enrichment);
72 }
73 self
74 }
75
76 #[must_use]
78 pub fn with_embedder(mut self, embedder: Arc<dyn Embedder + Send + Sync>) -> Self {
79 self.embedder = Some(embedder);
80 self
81 }
82
83 pub fn append(&self, record: &EpisodicRecord) -> Result<()> {
85 self.append_batch(std::slice::from_ref(record))
86 }
87
88 pub fn append_batch(&self, records: &[EpisodicRecord]) -> Result<()> {
91 if records.is_empty() {
92 return Ok(());
93 }
94 let serialized = records
95 .iter()
96 .map(|record| {
97 rmp_serde::to_vec(record)
98 .map(|value| (record, value))
99 .map_err(|e| CoreError::Memory(format!("episodic serialize failed: {e}")))
100 })
101 .collect::<Result<Vec<_>>>()?;
102 let mut tx = self
103 .env
104 .begin_rw_txn()
105 .map_err(|e| CoreError::Memory(format!("episodic rw_txn failed: {e}")))?;
106 for (record, value) in &serialized {
107 match tx.put(
108 self.db,
109 record.id.as_bytes(),
110 value,
111 WriteFlags::NO_OVERWRITE,
112 ) {
113 Ok(()) => {}
114 Err(lmdb::Error::KeyExist) => {
115 tx.abort();
116 return Err(CoreError::InvalidArgs(format!(
117 "episodic record {} already exists",
118 record.id
119 )));
120 }
121 Err(e) => {
122 tx.abort();
123 return Err(CoreError::Memory(format!("episodic append failed: {e}")));
124 }
125 }
126 }
127 tx.commit()
128 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
129 self.mutation_count
130 .fetch_add(records.len() as u64, std::sync::atomic::Ordering::Relaxed);
131 self.index_records(records)?;
135 self.clear_term_cache();
136 Ok(())
137 }
138
139 fn term_postings(&self, term: &str) -> Result<Vec<EpisodicId>> {
146 if let Ok(cache) = self.term_cache.read() {
147 if let Some(ids) = cache.get(term) {
148 return Ok(ids.clone());
149 }
150 }
151
152 let tx = self
153 .env
154 .begin_ro_txn()
155 .map_err(|e| CoreError::Memory(format!("episodic index ro_txn failed: {e}")))?;
156 let term_key = term.to_string();
157 let mut ids: Vec<EpisodicId> = Vec::new();
158 {
159 match tx.get(self.term_db, &term_key) {
162 Ok(_) => {}
163 Err(lmdb::Error::NotFound) => {
164 tx.commit().map_err(|e| {
165 CoreError::Memory(format!("episodic index commit failed: {e}"))
166 })?;
167 if let Ok(mut cache) = self.term_cache.write() {
168 cache.insert(term_key, Vec::new());
169 }
170 return Ok(ids);
171 }
172 Err(e) => {
173 return Err(CoreError::Memory(format!(
174 "episodic index read failed: {e}"
175 )));
176 }
177 }
178 let mut cursor = tx
179 .open_ro_cursor(self.term_db)
180 .map_err(|e| CoreError::Memory(format!("episodic index cursor failed: {e}")))?;
181 for (key, value) in cursor.iter_from(term_key.as_bytes()) {
182 if key != term_key.as_bytes() {
183 break;
184 }
185 if value.len() == std::mem::size_of::<EpisodicId>() {
186 if let Ok(id) = EpisodicId::from_slice(value) {
187 ids.push(id);
188 }
189 }
190 }
191 }
192 tx.commit()
193 .map_err(|e| CoreError::Memory(format!("episodic index commit failed: {e}")))?;
194 if let Ok(mut cache) = self.term_cache.write() {
195 cache.insert(term_key, ids.clone());
196 }
197 Ok(ids)
198 }
199
200 fn clear_term_cache(&self) {
201 if let Ok(mut cache) = self.term_cache.write() {
202 cache.clear();
203 }
204 }
205
206 fn index_records(&self, records: &[EpisodicRecord]) -> Result<()> {
207 let public: Vec<&EpisodicRecord> = records
208 .iter()
209 .filter(|record| !record.is_private && !record.model_exclude)
210 .collect();
211 if public.is_empty() {
212 return Ok(());
213 }
214 let mut pending: HashMap<String, Vec<&EpisodicRecord>> = HashMap::new();
215 for record in &public {
216 let base_terms = index_terms_with_aliases(&record.content, self.aliases.as_ref());
217 let enriched: Vec<String> = if let Some(ref enrichment) = self.enrichment {
218 let mut all = base_terms.clone();
219 let extra = enrichment.enrich(&base_terms);
220 all.extend(extra);
221 all.sort();
222 all.dedup();
223 all
224 } else {
225 base_terms
226 };
227 for term in enriched {
228 pending.entry(term).or_default().push(record);
229 }
230 }
231 let mut tx = self
232 .env
233 .begin_rw_txn()
234 .map_err(|e| CoreError::Memory(format!("episodic index rw_txn failed: {e}")))?;
235 for (term, records_for_term) in pending {
236 for record in records_for_term {
237 if let Err(e) = tx.put(
241 self.term_db,
242 &term.as_bytes().to_vec(),
243 &record.id.as_bytes(),
244 WriteFlags::default(),
245 ) {
246 tx.abort();
247 return Err(CoreError::Memory(format!(
248 "episodic term index write failed: {e}"
249 )));
250 }
251 }
252 }
253 tx.commit()
254 .map_err(|e| CoreError::Memory(format!("episodic term index commit failed: {e}")))?;
255 Ok(())
256 }
257
258 pub fn rebuild_sidecar(&self) -> Result<usize> {
265 let records = self.scan(None, usize::MAX)?;
266 let mut indexed = 0usize;
267 for chunk in records.chunks(5_000) {
268 self.index_records(chunk)?;
269 indexed += chunk.len();
270 }
271 self.clear_term_cache();
272 Ok(indexed)
273 }
274
275 pub fn sidecar_is_empty(&self) -> Result<bool> {
282 let tx = self
283 .env
284 .begin_ro_txn()
285 .map_err(|e| CoreError::Memory(format!("episodic index ro_txn failed: {e}")))?;
286 let mut cursor = tx
287 .open_ro_cursor(self.term_db)
288 .map_err(|e| CoreError::Memory(format!("episodic index cursor failed: {e}")))?;
289 Ok(cursor.iter().next().is_none())
290 }
291
292 pub fn record_count(&self) -> Result<u64> {
294 let tx = self
295 .env
296 .begin_ro_txn()
297 .map_err(|e| CoreError::Memory(format!("episodic ro_txn failed: {e}")))?;
298 let mut cursor = tx
299 .open_ro_cursor(self.db)
300 .map_err(|e| CoreError::Memory(format!("episodic cursor failed: {e}")))?;
301 let mut count = 0u64;
302 for _ in cursor.iter() {
303 count += 1;
304 }
305 Ok(count)
306 }
307
308 pub fn append_explicit(
310 &self,
311 record: &EpisodicRecord,
312 policy: EpisodicCapturePolicy,
313 ) -> Result<bool> {
314 let prepared = record
315 .clone()
316 .with_content(policy.prepare_content(&record.content));
317 self.append(&prepared)?;
318 Ok(true)
319 }
320
321 pub fn append_explicit_batch(
323 &self,
324 records: &[EpisodicRecord],
325 policy: EpisodicCapturePolicy,
326 ) -> Result<usize> {
327 if records.is_empty() {
328 return Ok(0);
329 }
330 let prepared: Vec<EpisodicRecord> = records
331 .iter()
332 .map(|record| {
333 record
334 .clone()
335 .with_content(policy.prepare_content(&record.content))
336 })
337 .collect();
338 self.append_batch(&prepared)?;
339 Ok(prepared.len())
340 }
341
342 pub fn get(&self, id: EpisodicId) -> Result<Option<EpisodicRecord>> {
344 let tx = self
345 .env
346 .begin_ro_txn()
347 .map_err(|e| CoreError::Memory(format!("episodic ro_txn failed: {e}")))?;
348 let result = tx.get(self.db, id.as_bytes());
349 match result {
350 Ok(bytes) => {
351 let record: EpisodicRecord = rmp_serde::from_slice(bytes)
352 .map_err(|e| CoreError::Memory(format!("episodic deserialize failed: {e}")))?;
353 tx.commit()
354 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
355 Ok(Some(record))
356 }
357 Err(lmdb::Error::NotFound) => {
358 tx.commit()
359 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
360 Ok(None)
361 }
362 Err(e) => Err(CoreError::Memory(format!("episodic get failed: {e}"))),
363 }
364 }
365
366 pub fn transition(&self, id: EpisodicId, transition: MemoryTransition) -> Result<()> {
368 let mut tx = self
369 .env
370 .begin_rw_txn()
371 .map_err(|e| CoreError::Memory(format!("episodic rw_txn failed: {e}")))?;
372 let bytes = match tx.get(self.db, id.as_bytes()) {
373 Ok(bytes) => bytes,
374 Err(lmdb::Error::NotFound) => {
375 tx.abort();
376 return Err(CoreError::InvalidArgs(format!(
377 "episodic record {id} does not exist"
378 )));
379 }
380 Err(e) => {
381 tx.abort();
382 return Err(CoreError::Memory(format!("episodic get failed: {e}")));
383 }
384 };
385 let mut record: EpisodicRecord = rmp_serde::from_slice(bytes)
386 .map_err(|e| CoreError::Memory(format!("episodic deserialize failed: {e}")))?;
387 record
388 .transition(transition)
389 .map_err(|e| CoreError::InvalidArgs(format!("episodic transition rejected: {e}")))?;
390 let value = rmp_serde::to_vec(&record)
391 .map_err(|e| CoreError::Memory(format!("episodic serialize failed: {e}")))?;
392 tx.put(self.db, id.as_bytes(), &value, WriteFlags::default())
393 .map_err(|e| CoreError::Memory(format!("episodic transition write failed: {e}")))?;
394 tx.commit()
395 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
396 self.mutation_count
397 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
398 Ok(())
399 }
400
401 pub fn scan(
403 &self,
404 session_id: Option<uuid::Uuid>,
405 limit: usize,
406 ) -> Result<Vec<EpisodicRecord>> {
407 if limit == 0 {
408 return Ok(Vec::new());
409 }
410 let tx = self
411 .env
412 .begin_ro_txn()
413 .map_err(|e| CoreError::Memory(format!("episodic ro_txn failed: {e}")))?;
414 let mut cursor = tx
415 .open_ro_cursor(self.db)
416 .map_err(|e| CoreError::Memory(format!("episodic cursor failed: {e}")))?;
417 let mut records = Vec::new();
418 for item in cursor.iter() {
419 let (_, bytes) = item;
420 let record: EpisodicRecord = rmp_serde::from_slice(bytes)
421 .map_err(|e| CoreError::Memory(format!("episodic deserialize failed: {e}")))?;
422 if session_id.is_none_or(|id| record.session_id == Some(id)) {
423 records.push(record);
424 }
425 }
426 drop(cursor);
427 tx.commit()
428 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
429 records.sort_by_key(|record| (record.sequence, record.created_at, record.id));
430 records.truncate(limit);
431 Ok(records)
432 }
433
434 pub fn search(
444 &self,
445 query: &str,
446 limit: usize,
447 include_historical: bool,
448 ) -> Result<Vec<EpisodicSearchResult>> {
449 self.search_with_limits(query, limit, limit.saturating_mul(2), include_historical)
450 }
451
452 pub fn search_with_limits(
460 &self,
461 query: &str,
462 limit: usize,
463 candidate_limit: usize,
464 include_historical: bool,
465 ) -> Result<Vec<EpisodicSearchResult>> {
466 if limit == 0 {
467 return Ok(Vec::new());
468 }
469 let mut results = self.search_scored(query, limit, candidate_limit, include_historical)?;
470 if is_current_query(query) {
471 self.resolve_current(&mut results);
472 }
473 results.truncate(limit);
474 Ok(results)
475 }
476
477 fn search_scored(
482 &self,
483 query: &str,
484 limit: usize,
485 candidate_limit: usize,
486 include_historical: bool,
487 ) -> Result<Vec<EpisodicSearchResult>> {
488 let plan = QueryPlan::plan(query, limit);
489 let candidate_limit = candidate_limit.max(plan.candidate_limit);
490 let query_terms = tokenize(query);
491 let query_keys = key_index_terms_with_aliases(query, self.aliases.as_ref());
492 if query_terms.is_empty() && query_keys.is_empty() || limit == 0 {
493 return Ok(Vec::new());
494 }
495 let mut candidate_scores: HashMap<EpisodicId, usize> = HashMap::new();
496 for term in query_terms.iter().chain(query_keys.iter()) {
497 for id in self.term_postings(term)? {
498 *candidate_scores.entry(id).or_default() += 1;
499 }
500 }
501
502 let records = if candidate_scores.is_empty() {
503 if !matches!(self.sidecar_is_empty(), Ok(true)) {
508 return Ok(Vec::new());
509 }
510 self.scan(None, usize::MAX)?
512 } else {
513 let mut ranked_candidates: Vec<(EpisodicId, usize)> =
514 candidate_scores.into_iter().collect();
515 ranked_candidates.sort_by(|(left_id, left_count), (right_id, right_count)| {
516 right_count
517 .cmp(left_count)
518 .then_with(|| left_id.cmp(right_id))
519 });
520 ranked_candidates.truncate(candidate_limit);
521 self.load_records(
522 &ranked_candidates
523 .into_iter()
524 .map(|(id, _)| id)
525 .collect::<Vec<_>>(),
526 )?
527 };
528
529 let mut results = Vec::new();
530 for record in records {
531 if !include_historical && !matches!(record.validity, ValidityState::Active) {
532 continue;
533 }
534 let content_terms = tokenize(&record.content);
535 let content_keys = key_index_terms_with_aliases(&record.content, self.aliases.as_ref());
536 let reverse_map: HashMap<&String, Vec<String>> =
541 if let Some(ref enrichment) = self.enrichment {
542 if matches!(record.kind, EpisodicKind::UserStatement) {
543 query_terms
544 .iter()
545 .map(|qt| (qt, enrichment.reverse_enrich(qt)))
546 .collect()
547 } else {
548 HashMap::new()
549 }
550 } else {
551 HashMap::new()
552 };
553 let mut reverse_match_count = 0usize;
554 let matched_terms = query_terms
555 .iter()
556 .filter(|term| {
557 if content_terms.iter().any(|candidate| candidate == *term)
558 || content_keys.iter().any(|candidate| candidate == *term)
559 {
560 return true;
561 }
562 if let Some(reverse_terms) = reverse_map.get(term) {
565 let found = reverse_terms.iter().any(|rt| {
566 content_terms.iter().any(|candidate| candidate == rt)
567 || content_keys.iter().any(|candidate| candidate == rt)
568 });
569 if found {
570 reverse_match_count += 1;
571 }
572 return found;
573 }
574 false
575 })
576 .count();
577 let matched_keys = query_keys
578 .iter()
579 .filter(|term| {
580 content_keys.iter().any(|candidate| candidate == *term)
581 || content_terms.iter().any(|candidate| candidate == *term)
582 })
583 .count();
584 if matched_terms == 0 && matched_keys == 0 {
585 continue;
586 }
587 let key_bonus = if query_keys.is_empty() {
588 0.0
589 } else {
590 matched_keys as f32 / query_keys.len() as f32 * plan.key_weight
591 };
592 let role_boost = match record.kind {
593 EpisodicKind::UserStatement => 0.12,
594 _ => 0.0,
595 };
596 let effective_matched = if matches!(record.kind, EpisodicKind::UserStatement) {
597 (matched_terms + 2).min(query_terms.len())
598 } else {
599 matched_terms
600 };
601 let coverage = if query_terms.is_empty() {
602 0.0
603 } else {
604 effective_matched as f32 / query_terms.len() as f32
605 };
606 let number_bonus = if plan.number_query {
607 let has_digit = content_terms
608 .iter()
609 .any(|term| term.chars().any(|c| c.is_ascii_digit()));
610 if has_digit || contains_number_word(&record.content) {
611 0.03
612 } else {
613 0.0
614 }
615 } else {
616 0.0
617 };
618 let density = matched_terms as f32 / content_terms.len().max(1) as f32;
619 results.push(EpisodicSearchResult {
620 record,
621 score: coverage
622 + key_bonus
623 + role_boost
624 + number_bonus
625 + (reverse_match_count as f32).mul_add(0.05, density * 0.03),
626 matched_terms: matched_terms.max(matched_keys),
627 });
628 }
629 let mut session_counts: HashMap<Option<uuid::Uuid>, usize> = HashMap::new();
633 for r in &results {
634 *session_counts.entry(r.record.session_id).or_default() += 1;
635 }
636 for r in &mut results {
637 let count = session_counts
638 .get(&r.record.session_id)
639 .copied()
640 .unwrap_or(1);
641 if count > 1 {
642 r.score = 0.02f32.mul_add((count - 1).min(3) as f32, r.score);
643 }
644 }
645 let mut hash_counts: HashMap<&str, usize> = HashMap::new();
651 for r in &results {
652 *hash_counts
653 .entry(r.record.content_hash.as_str())
654 .or_default() += 1;
655 }
656 let hash_boosts: HashMap<String, f32> = results
657 .iter()
658 .map(|r| {
659 let count = hash_counts
660 .get(r.record.content_hash.as_str())
661 .copied()
662 .unwrap_or(1);
663 let boost = if count > 1 {
664 0.03 * (count - 1).min(3) as f32
665 } else {
666 0.0
667 };
668 (r.record.id.to_string(), boost)
669 })
670 .collect();
671 for r in &mut results {
672 if let Some(boost) = hash_boosts.get(&r.record.id.to_string()) {
673 r.score += boost;
674 }
675 }
676 results.sort_by(|a, b| {
677 b.score
678 .partial_cmp(&a.score)
679 .unwrap_or(std::cmp::Ordering::Equal)
680 .then_with(|| b.matched_terms.cmp(&a.matched_terms))
681 .then_with(|| a.record.content.len().cmp(&b.record.content.len()))
682 .then_with(|| a.record.sequence.cmp(&b.record.sequence))
683 .then_with(|| a.record.id.cmp(&b.record.id))
684 });
685 Ok(results)
686 }
687
688 pub fn search_with_rerank(
705 &self,
706 query: &str,
707 limit: usize,
708 candidate_limit: usize,
709 include_historical: bool,
710 alpha: f32,
711 ) -> Result<Vec<EpisodicSearchResult>> {
712 let Some(ref embedder) = self.embedder else {
713 return self.search_with_limits(query, limit, candidate_limit, include_historical);
714 };
715 if !embedder.is_available() || limit == 0 {
716 return self.search_with_limits(query, limit, candidate_limit, include_historical);
717 }
718
719 let rerank_pool = limit.max(candidate_limit).min(50);
721 let deterministic =
722 self.search_scored(query, rerank_pool, rerank_pool, include_historical)?;
723 if deterministic.is_empty() {
724 return Ok(Vec::new());
725 }
726
727 let contents: Vec<&str> = std::iter::once(query)
729 .chain(deterministic.iter().map(|r| r.record.content.as_str()))
730 .collect();
731 let embeddings = embedder.embed_batch(&contents)?;
732 if embeddings.len() != deterministic.len() + 1 {
733 return Err(CoreError::Memory(format!(
734 "embedder returned {} vectors, expected {}",
735 embeddings.len(),
736 deterministic.len() + 1
737 )));
738 }
739 let query_vec = &embeddings[0];
740 let candidate_vecs = &embeddings[1..];
741
742 if alpha >= 2.0 {
743 let protected: Vec<EpisodicSearchResult> =
748 deterministic.into_iter().take(limit).collect();
749 let cosines: Vec<f32> = protected
750 .iter()
751 .enumerate()
752 .map(|(i, _)| cosine_sim(query_vec, &candidate_vecs[i]))
753 .collect();
754 let mut order: Vec<usize> = (0..protected.len()).collect();
755 order.sort_by(|&a, &b| {
756 cosines[b]
757 .partial_cmp(&cosines[a])
758 .unwrap_or(std::cmp::Ordering::Equal)
759 .then_with(|| a.cmp(&b))
761 });
762 let mut slots: Vec<Option<EpisodicSearchResult>> =
763 protected.into_iter().map(Some).collect();
764 let reranked: Vec<EpisodicSearchResult> =
765 order.into_iter().filter_map(|i| slots[i].take()).collect();
766 Ok(reranked)
767 } else if alpha >= 1.0 {
768 let delta = 0.05;
770 let mut reranked = deterministic;
771 let cosines: Vec<f32> = candidate_vecs
772 .iter()
773 .map(|v| cosine_sim(query_vec, v))
774 .collect();
775 let n = reranked.len();
777 for _ in 0..n {
778 let mut swapped = false;
779 for i in 0..n.saturating_sub(1) {
780 let det_gap = (reranked[i].score - reranked[i + 1].score).abs();
781 if det_gap < delta && cosines[i + 1] > cosines[i] {
782 reranked.swap(i, i + 1);
783 swapped = true;
784 }
785 }
786 if !swapped {
787 break;
788 }
789 }
790 if is_current_query(query) {
791 self.resolve_current(&mut reranked);
792 }
793 reranked.truncate(limit);
794 Ok(reranked)
795 } else {
796 let max_det = deterministic
798 .iter()
799 .map(|r| r.score)
800 .fold(0.0f32, f32::max)
801 .max(1e-9);
802
803 let mut reranked: Vec<EpisodicSearchResult> = deterministic
804 .into_iter()
805 .enumerate()
806 .map(|(i, mut r)| {
807 let cosine = cosine_sim(query_vec, &candidate_vecs[i]);
808 let det_norm = r.score / max_det;
809 r.score = alpha.mul_add(det_norm, (1.0 - alpha) * cosine);
810 r
811 })
812 .collect();
813
814 reranked.sort_by(|a, b| {
815 b.score
816 .partial_cmp(&a.score)
817 .unwrap_or(std::cmp::Ordering::Equal)
818 .then_with(|| b.matched_terms.cmp(&a.matched_terms))
819 .then_with(|| a.record.content.len().cmp(&b.record.content.len()))
820 .then_with(|| a.record.sequence.cmp(&b.record.sequence))
821 .then_with(|| a.record.id.cmp(&b.record.id))
822 });
823 if is_current_query(query) {
824 self.resolve_current(&mut reranked);
825 }
826 reranked.truncate(limit);
827 Ok(reranked)
828 }
829 }
830
831 fn resolve_current(&self, results: &mut Vec<EpisodicSearchResult>) {
858 if results.len() < 2 {
859 return;
860 }
861 let mut anchors: Vec<EpisodicSearchResult> = Vec::new();
862 let mut rest: Vec<EpisodicSearchResult> = Vec::new();
863 for result in results.drain(..) {
864 let is_anchor = matches!(result.record.kind, EpisodicKind::UserStatement)
865 && contains_change_marker(&result.record.content);
866 if is_anchor {
867 anchors.push(result);
868 } else {
869 rest.push(result);
870 }
871 }
872 if anchors.is_empty() {
873 *results = rest;
875 return;
876 }
877 anchors.sort_by(|a, b| {
878 b.record
879 .created_at
880 .cmp(&a.record.created_at)
881 .then_with(|| b.record.sequence.cmp(&a.record.sequence))
882 .then_with(|| a.record.id.cmp(&b.record.id))
883 });
884 anchors.extend(rest);
885 *results = anchors;
886 }
887
888 fn load_records(&self, ids: &[EpisodicId]) -> Result<Vec<EpisodicRecord>> {
889 let tx = self
890 .env
891 .begin_ro_txn()
892 .map_err(|e| CoreError::Memory(format!("episodic ro_txn failed: {e}")))?;
893 let mut records = Vec::with_capacity(ids.len());
894 for id in ids {
895 match tx.get(self.db, id.as_bytes()) {
896 Ok(bytes) => {
897 records.push(rmp_serde::from_slice(bytes).map_err(|e| {
898 CoreError::Memory(format!("episodic deserialize failed: {e}"))
899 })?);
900 }
901 Err(lmdb::Error::NotFound) => {}
902 Err(e) => return Err(CoreError::Memory(format!("episodic get failed: {e}"))),
903 }
904 }
905 tx.commit()
906 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
907 Ok(records)
908 }
909}
910
911fn index_terms_with_aliases(text: &str, aliases: Option<&AdaptiveAliases>) -> Vec<String> {
912 tokenize(text)
913 .into_iter()
914 .chain(key_index_terms_with_aliases(text, aliases))
915 .fold(Vec::new(), |mut terms, term| {
916 if !terms.contains(&term) {
917 terms.push(term);
918 }
919 terms
920 })
921}
922
923fn tokenize(text: &str) -> Vec<String> {
924 strip_stopwords(text)
925 .split(|c: char| !c.is_alphanumeric())
926 .filter(|term| term.len() > 1)
927 .map(|term| simple_stem(&term.to_ascii_lowercase()))
928 .fold(Vec::new(), |mut terms, term| {
929 if !terms.contains(&term) {
930 terms.push(term);
931 }
932 terms
933 })
934}
935
936const CURRENT_QUERY_WORD_CUES: &[&str] = &["current", "currently", "latest", "nowadays"];
938
939const CURRENT_QUERY_PHRASE_CUES: &[&str] = &["these days", "right now", "at the moment"];
941
942#[must_use]
949pub fn is_current_query(query: &str) -> bool {
950 let lowered = query.to_ascii_lowercase();
951 let has_word = lowered
952 .split(|c: char| !c.is_alphanumeric())
953 .any(|token| CURRENT_QUERY_WORD_CUES.contains(&token));
954 has_word
955 || CURRENT_QUERY_PHRASE_CUES
956 .iter()
957 .any(|cue| lowered.contains(cue))
958}
959
960const CHANGE_MARKERS: &[&str] = &[
965 "switched to",
966 "switch to",
967 "switching to",
968 "switched from",
969 "changed my",
970 "change my",
971 "changed from",
972 "now prefer",
973 "now i prefer",
974 "now i'm",
975 "now im",
976 "no longer",
977 "used to",
978 "moved to",
979 "not anymore",
980 "instead of",
981 "replaced",
982 "gave up",
983];
984
985fn contains_change_marker(content: &str) -> bool {
987 let lowered = content.to_ascii_lowercase();
988 CHANGE_MARKERS.iter().any(|marker| lowered.contains(marker))
989}
990
991const CONTRADICTION_MARKERS: &[&str] = &[
995 "no longer",
996 "anymore",
997 "changed my mind",
998 "changed my",
999 "used to",
1000 "gave up",
1001 "just a phase",
1002 "not really",
1003 "but i",
1004];
1005
1006#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1013pub struct EpisodicConflict {
1014 pub later_record: EpisodicId,
1017 pub earlier_record: EpisodicId,
1019 pub marker: String,
1021 pub shared_terms: Vec<String>,
1023 pub later_content: String,
1025 pub earlier_content: String,
1027}
1028
1029#[must_use]
1037pub fn detect_conflicts(results: &[EpisodicSearchResult]) -> Vec<EpisodicConflict> {
1038 const MAX_CONFLICTS: usize = 10;
1039 let mut conflicts: Vec<EpisodicConflict> = Vec::new();
1040 let mut seen_pairs: Vec<(EpisodicId, EpisodicId)> = Vec::new();
1041 for (i, marked) in results.iter().enumerate() {
1042 if !matches!(marked.record.kind, EpisodicKind::UserStatement) {
1043 continue;
1044 }
1045 let lowered = marked.record.content.to_ascii_lowercase();
1046 let Some(marker) = CONTRADICTION_MARKERS
1047 .iter()
1048 .find(|m| lowered.contains(*m))
1049 .copied()
1050 else {
1051 continue;
1052 };
1053 let marked_terms = tokenize(&marked.record.content);
1054 for (j, other) in results.iter().enumerate() {
1055 if i == j || !matches!(other.record.kind, EpisodicKind::UserStatement) {
1056 continue;
1057 }
1058 if other.record.content_hash == marked.record.content_hash {
1059 continue;
1060 }
1061 let other_terms = tokenize(&other.record.content);
1062 let shared: Vec<String> = marked_terms
1063 .iter()
1064 .filter(|t| other_terms.contains(t))
1065 .cloned()
1066 .collect();
1067 if shared.len() < 2 {
1068 continue;
1069 }
1070 let (later, earlier) = if (marked.record.created_at, marked.record.sequence)
1073 > (other.record.created_at, other.record.sequence)
1074 {
1075 (&marked.record, &other.record)
1076 } else {
1077 (&other.record, &marked.record)
1078 };
1079 let pair_key = if later.id < earlier.id {
1080 (later.id, earlier.id)
1081 } else {
1082 (earlier.id, later.id)
1083 };
1084 if seen_pairs.contains(&pair_key) {
1085 continue;
1086 }
1087 seen_pairs.push(pair_key);
1088 conflicts.push(EpisodicConflict {
1089 later_record: later.id,
1090 earlier_record: earlier.id,
1091 marker: marker.to_string(),
1092 shared_terms: shared,
1093 later_content: later.content.clone(),
1094 earlier_content: earlier.content.clone(),
1095 });
1096 if conflicts.len() >= MAX_CONFLICTS {
1097 return conflicts;
1098 }
1099 }
1100 }
1101 conflicts
1102}
1103
1104fn simple_stem(word: &str) -> String {
1105 if word.len() <= 3 {
1106 return word.to_string();
1107 }
1108 for suffix in ["ies", "ied", "ing", "edly", "ed", "ly", "es", "s"] {
1109 if let Some(stem) = word.strip_suffix(suffix) {
1110 if suffix == "ies" || suffix == "ied" {
1111 return format!("{stem}y");
1112 }
1113 if stem.len() >= 2 {
1114 return stem.to_string();
1115 }
1116 }
1117 }
1118 word.to_string()
1119}
1120
1121fn cosine_sim(a: &[f32], b: &[f32]) -> f32 {
1122 let dot = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum::<f32>();
1123 let norm_a = a.iter().map(|x| x * x).sum::<f32>().sqrt();
1124 let norm_b = b.iter().map(|x| x * x).sum::<f32>().sqrt();
1125 if norm_a < 1e-9 || norm_b < 1e-9 {
1126 0.0
1127 } else {
1128 dot / (norm_a * norm_b)
1129 }
1130}
1131
1132fn contains_number_word(text: &str) -> bool {
1133 const NUMBER_WORDS: &[&str] = &[
1134 "one",
1135 "two",
1136 "three",
1137 "four",
1138 "five",
1139 "six",
1140 "seven",
1141 "eight",
1142 "nine",
1143 "ten",
1144 "eleven",
1145 "twelve",
1146 "thirteen",
1147 "fourteen",
1148 "fifteen",
1149 "sixteen",
1150 "seventeen",
1151 "eighteen",
1152 "nineteen",
1153 "twenty",
1154 "thirty",
1155 "forty",
1156 "fifty",
1157 "sixty",
1158 "seventy",
1159 "eighty",
1160 "ninety",
1161 "hundred",
1162 "thousand",
1163 "million",
1164 "billion",
1165 "dozen",
1166 "couple",
1167 "half",
1168 "quarter",
1169 "double",
1170 "triple",
1171 "twice",
1172 ];
1173 for word in text.split(|c: char| !c.is_alphanumeric()) {
1174 if word.len() >= 3 && NUMBER_WORDS.iter().any(|nw| word.eq_ignore_ascii_case(nw)) {
1175 return true;
1176 }
1177 }
1178 false
1179}
1180
1181#[cfg(test)]
1182mod tests {
1183 use super::*;
1184 use crate::MemoryStore;
1185 use tempfile::tempdir;
1186 use wm_core::{EpisodicKind, Provenance, ProvenanceSource, ValidityState};
1187
1188 fn sample_record(sequence: u64, content: &str) -> EpisodicRecord {
1189 EpisodicRecord::new(
1190 None,
1191 sequence,
1192 EpisodicKind::Observation,
1193 content,
1194 Provenance::new(ProvenanceSource::User),
1195 )
1196 }
1197
1198 fn user_statement(sequence: u64, content: &str) -> EpisodicRecord {
1199 EpisodicRecord::new(
1200 None,
1201 sequence,
1202 EpisodicKind::UserStatement,
1203 content,
1204 Provenance::new(ProvenanceSource::User),
1205 )
1206 }
1207
1208 fn assistant_response(sequence: u64, content: &str) -> EpisodicRecord {
1209 EpisodicRecord::new(
1210 None,
1211 sequence,
1212 EpisodicKind::AssistantResponse,
1213 content,
1214 Provenance::new(ProvenanceSource::Agent),
1215 )
1216 }
1217
1218 #[test]
1219 fn current_query_detection() {
1220 assert!(is_current_query("What's my current favorite coffee?"));
1221 assert!(is_current_query("What am I currently reading these days?"));
1222 assert!(is_current_query("What's the latest book I mentioned?"));
1223 assert!(is_current_query("What's my job right now?"));
1224 assert!(is_current_query("What am I eating at the moment?"));
1225 assert!(!is_current_query("What's my favorite coffee?"));
1226 assert!(!is_current_query("Where did I volunteer in February?"));
1227 assert!(!is_current_query("What did I say about the trip?"));
1228 assert!(!is_current_query("What currency did I use in Japan?"));
1230 }
1231
1232 #[test]
1233 fn current_query_resolution_prefers_latest_statement() {
1234 let tmp = tempdir().unwrap();
1235 let store = MemoryStore::open_default(tmp.path()).unwrap();
1236 let episodic = store.episodic();
1237 episodic
1241 .append(&user_statement(1, "My favorite coffee is dark roast."))
1242 .unwrap();
1243 episodic
1244 .append(&user_statement(
1245 2,
1246 "I really love dark roast when it comes to coffee.",
1247 ))
1248 .unwrap();
1249 episodic
1250 .append(&user_statement(3, "I've been jogging lately."))
1251 .unwrap();
1252 episodic
1253 .append(&user_statement(4, "I've switched to cold brew for coffee."))
1254 .unwrap();
1255
1256 let results = episodic
1257 .search("What's my current favorite coffee?", 5, false)
1258 .unwrap();
1259 assert!(!results.is_empty());
1260 assert!(
1261 results[0].record.content.contains("cold brew"),
1262 "current query must rank the latest statement first, got: {}",
1263 results[0].record.content
1264 );
1265 }
1266
1267 #[test]
1268 fn current_query_anchors_switched_from_template() {
1269 let tmp = tempdir().unwrap();
1275 let store = MemoryStore::open_default(tmp.path()).unwrap();
1276 let episodic = store.episodic();
1277 episodic
1278 .append(&user_statement(1, "My favorite coffee is espresso."))
1279 .unwrap();
1280 episodic
1281 .append(&user_statement(
1282 2,
1283 "I've actually switched from espresso to latte for coffee.",
1284 ))
1285 .unwrap();
1286 episodic
1287 .append(&user_statement(3, "My favorite coffee is latte."))
1288 .unwrap();
1289
1290 let results = episodic
1291 .search("What's my current favorite coffee?", 5, false)
1292 .unwrap();
1293 assert!(!results.is_empty());
1294 assert!(
1295 results[0].record.content.contains("latte"),
1296 "'switched from' must anchor the current value, got: {}",
1297 results[0].record.content
1298 );
1299 }
1300
1301 #[test]
1302 fn non_current_query_keeps_score_order() {
1303 let tmp = tempdir().unwrap();
1304 let store = MemoryStore::open_default(tmp.path()).unwrap();
1305 let episodic = store.episodic();
1306 episodic
1307 .append(&user_statement(1, "My favorite coffee is dark roast."))
1308 .unwrap();
1309 episodic
1310 .append(&user_statement(2, "I've switched to cold brew for coffee."))
1311 .unwrap();
1312
1313 let results = episodic
1316 .search("What's my favorite coffee?", 5, false)
1317 .unwrap();
1318 assert!(!results.is_empty());
1319 assert!(
1320 results[0].record.content.contains("dark roast"),
1321 "non-current query must keep score order, got: {}",
1322 results[0].record.content
1323 );
1324 }
1325
1326 #[test]
1327 fn current_resolution_anchors_on_user_statements_only() {
1328 let tmp = tempdir().unwrap();
1329 let store = MemoryStore::open_default(tmp.path()).unwrap();
1330 let episodic = store.episodic();
1331 episodic
1332 .append(&user_statement(1, "My favorite coffee is dark roast."))
1333 .unwrap();
1334 episodic
1336 .append(&assistant_response(
1337 2,
1338 "Got it, dark roast is your favorite coffee!",
1339 ))
1340 .unwrap();
1341 episodic
1342 .append(&user_statement(3, "I've switched to cold brew for coffee."))
1343 .unwrap();
1344
1345 let results = episodic
1346 .search("What's my current favorite coffee?", 5, false)
1347 .unwrap();
1348 assert!(
1349 results[0].record.content.contains("cold brew"),
1350 "user statements anchor chronology, got: {}",
1351 results[0].record.content
1352 );
1353 assert_eq!(results[0].record.kind, EpisodicKind::UserStatement);
1354 }
1355
1356 #[test]
1357 fn current_query_without_change_markers_keeps_score_order() {
1358 let tmp = tempdir().unwrap();
1359 let store = MemoryStore::open_default(tmp.path()).unwrap();
1360 let episodic = store.episodic();
1361 episodic
1363 .append(&user_statement(
1364 1,
1365 "My favorite hiking trail is Eagle Ridge.",
1366 ))
1367 .unwrap();
1368 episodic
1369 .append(&user_statement(2, "I go hiking every weekend."))
1370 .unwrap();
1371
1372 let results = episodic
1373 .search("What's my current favorite hiking trail?", 5, false)
1374 .unwrap();
1375 assert!(!results.is_empty());
1376 assert!(
1377 results[0].record.content.contains("Eagle Ridge"),
1378 "no change markers → deterministic score order, got: {}",
1379 results[0].record.content
1380 );
1381 }
1382
1383 #[test]
1384 fn detect_conflicts_flags_contradiction_with_shared_topic() {
1385 let tmp = tempdir().unwrap();
1386 let store = MemoryStore::open_default(tmp.path()).unwrap();
1387 let episodic = store.episodic();
1388 episodic
1389 .append(&user_statement(
1390 1,
1391 "I'm vegetarian now. I decided to stop eating animal products.",
1392 ))
1393 .unwrap();
1394 episodic
1395 .append(&user_statement(
1396 2,
1397 "I'm not really vegetarian anymore, I eat steak now.",
1398 ))
1399 .unwrap();
1400 episodic
1401 .append(&user_statement(3, "I went hiking yesterday."))
1402 .unwrap();
1403
1404 let results = episodic
1405 .search("vegetarian steak eating", 10, false)
1406 .unwrap();
1407 let conflicts = detect_conflicts(&results);
1408 assert_eq!(
1409 conflicts.len(),
1410 1,
1411 "the vegetarian/steak pair must be flagged, got {conflicts:?}"
1412 );
1413 let conflict = &conflicts[0];
1414 assert!(conflict.later_content.contains("steak"));
1415 assert!(conflict.earlier_content.contains("animal products"));
1416 assert!(
1417 conflict.shared_terms.iter().any(|t| t == "vegetarian"),
1418 "shared terms must include the topic: {:?}",
1419 conflict.shared_terms
1420 );
1421 }
1422
1423 #[test]
1424 fn detect_conflicts_ignores_plain_statements_and_assistant_turns() {
1425 let tmp = tempdir().unwrap();
1426 let store = MemoryStore::open_default(tmp.path()).unwrap();
1427 let episodic = store.episodic();
1428 episodic
1430 .append(&user_statement(1, "My favorite coffee is dark roast."))
1431 .unwrap();
1432 episodic
1433 .append(&user_statement(2, "I love coffee with breakfast."))
1434 .unwrap();
1435 episodic
1437 .append(&assistant_response(
1438 3,
1439 "You mentioned you no longer like tea!",
1440 ))
1441 .unwrap();
1442
1443 let results = episodic.search("coffee tea breakfast", 10, false).unwrap();
1444 assert!(detect_conflicts(&results).is_empty());
1445 }
1446
1447 #[test]
1448 fn detect_conflicts_skips_identical_content() {
1449 let tmp = tempdir().unwrap();
1450 let store = MemoryStore::open_default(tmp.path()).unwrap();
1451 let episodic = store.episodic();
1452 let record = user_statement(1, "I'm vegetarian now, but I changed my mind.");
1453 let duplicate = user_statement(2, "I'm vegetarian now, but I changed my mind.");
1454 episodic.append(&record).unwrap();
1455 episodic.append(&duplicate).unwrap();
1456
1457 let results = episodic.search("vegetarian", 10, false).unwrap();
1458 assert!(detect_conflicts(&results).is_empty());
1460 }
1461
1462 #[test]
1463 fn append_get_transition_and_reopen_roundtrip() {
1464 let tmp = tempdir().unwrap();
1465 let session = uuid::Uuid::new_v4();
1466 let record = EpisodicRecord::new(
1467 Some(session),
1468 2,
1469 EpisodicKind::Decision,
1470 "use the raw episodic lane",
1471 Provenance::new(ProvenanceSource::User).with_actor("test"),
1472 );
1473 let id = record.id;
1474 {
1475 let store = MemoryStore::open_default(tmp.path()).unwrap();
1476 let episodic = store.episodic();
1477 episodic.append(&record).unwrap();
1478 assert_eq!(episodic.get(id).unwrap().unwrap(), record);
1479 episodic
1480 .transition(
1481 id,
1482 MemoryTransition::Supersede {
1483 replacement: uuid::Uuid::new_v4(),
1484 },
1485 )
1486 .unwrap();
1487 assert!(matches!(
1488 episodic.get(id).unwrap().unwrap().validity,
1489 ValidityState::Superseded { .. }
1490 ));
1491 }
1492 let reopened = MemoryStore::open_default(tmp.path()).unwrap();
1493 let records = reopened.episodic().scan(Some(session), 10).unwrap();
1494 assert_eq!(records.len(), 1);
1495 assert_eq!(records[0].id, id);
1496 }
1497
1498 #[test]
1499 fn duplicate_append_is_rejected() {
1500 let tmp = tempdir().unwrap();
1501 let store = MemoryStore::open_default(tmp.path()).unwrap();
1502 let record = sample_record(1, "once");
1503 store.episodic().append(&record).unwrap();
1504 let error = store.episodic().append(&record).unwrap_err();
1505 assert!(error.to_string().contains("already exists"));
1506 }
1507
1508 #[test]
1509 fn raw_search_returns_canonical_records_and_skips_revoked_by_default() {
1510 let tmp = tempdir().unwrap();
1511 let store = MemoryStore::open_default(tmp.path()).unwrap();
1512 let active = sample_record(1, "Rust memory retrieval");
1513 let revoked = sample_record(2, "Rust memory retrieval old");
1514 let revoked_id = revoked.id;
1515 store.episodic().append(&active).unwrap();
1516 store.episodic().append(&revoked).unwrap();
1517 store
1518 .episodic()
1519 .transition(
1520 revoked_id,
1521 MemoryTransition::Revoke {
1522 reason: "stale".into(),
1523 },
1524 )
1525 .unwrap();
1526
1527 let current = store
1528 .episodic()
1529 .search("memory retrieval", 10, false)
1530 .unwrap();
1531 assert_eq!(current.len(), 1);
1532 assert_eq!(current[0].record.id, active.id);
1533
1534 let all = store
1535 .episodic()
1536 .search("memory retrieval", 10, true)
1537 .unwrap();
1538 assert_eq!(all.len(), 2);
1539 }
1540
1541 #[test]
1542 fn append_batch_indexes_once_and_preserves_search() {
1543 let tmp = tempdir().unwrap();
1544 let store = MemoryStore::open_default(tmp.path()).unwrap();
1545 let first = sample_record(1, "Dr. Patel scheduled a follow-up appointment");
1546 let second = sample_record(2, "unrelated grocery list");
1547 let first_id = first.id;
1548 store.episodic().append_batch(&[first, second]).unwrap();
1549 let hits = store
1550 .episodic()
1551 .search("patel appointment", 10, false)
1552 .unwrap();
1553 assert_eq!(hits.len(), 1);
1554 assert_eq!(hits[0].record.id, first_id);
1555 }
1556
1557 #[test]
1558 fn rejected_append_batch_preserves_prior_state_after_reopen() {
1559 let tmp = tempdir().unwrap();
1563 let original = sample_record(1, "acknowledged original record");
1564 let rejected_new = sample_record(2, "must not survive rejected batch");
1565 {
1566 let store = MemoryStore::open_default(tmp.path()).unwrap();
1567 store.episodic().append(&original).unwrap();
1568 let error = store
1569 .episodic()
1570 .append_batch(&[rejected_new.clone(), original.clone()])
1571 .unwrap_err();
1572 assert!(error.to_string().contains("already exists"));
1573 assert_eq!(
1574 store.episodic().get(original.id).unwrap(),
1575 Some(original.clone())
1576 );
1577 assert_eq!(store.episodic().get(rejected_new.id).unwrap(), None);
1578 }
1579
1580 let reopened = MemoryStore::open_default(tmp.path()).unwrap();
1581 assert_eq!(
1582 reopened.episodic().get(original.id).unwrap(),
1583 Some(original)
1584 );
1585 assert_eq!(reopened.episodic().get(rejected_new.id).unwrap(), None);
1586 }
1587
1588 #[test]
1589 fn lost_episodic_sidecar_rebuilds_from_raw_records_after_reopen() {
1590 let tmp = tempdir().unwrap();
1595 let record = sample_record(1, "sidecar recovery preserves searchable evidence");
1596 {
1597 let store = MemoryStore::open_default(tmp.path()).unwrap();
1598 let episodic = store.episodic();
1599 episodic.append(&record).unwrap();
1600 assert!(!episodic.sidecar_is_empty().unwrap());
1601
1602 let mut tx = store.env().begin_rw_txn().unwrap();
1603 tx.clear_db(episodic.term_db).unwrap();
1604 tx.commit().unwrap();
1605 assert!(episodic.sidecar_is_empty().unwrap());
1606 }
1607
1608 let reopened = MemoryStore::open_default(tmp.path()).unwrap();
1609 let episodic = reopened.episodic();
1610 assert_eq!(episodic.get(record.id).unwrap(), Some(record.clone()));
1611 assert!(!episodic.sidecar_is_empty().unwrap());
1612 let hits = episodic
1613 .search("sidecar searchable evidence", 10, false)
1614 .unwrap();
1615 assert_eq!(hits.len(), 1);
1616 assert_eq!(hits[0].record.id, record.id);
1617 }
1618
1619 #[test]
1620 fn append_explicit_batch_redacts_and_skips_private() {
1621 let tmp = tempdir().unwrap();
1622 let store = MemoryStore::open_default(tmp.path()).unwrap();
1623 let public = sample_record(1, "api_key=supersecret rust retrieval");
1624 let private = sample_record(2, "private rust retrieval").with_visibility(true, false);
1625 let public_id = public.id;
1626 store
1627 .episodic()
1628 .append_explicit_batch(&[public, private], EpisodicCapturePolicy::explicit_only())
1629 .unwrap();
1630 let stored = store.episodic().get(public_id).unwrap().unwrap();
1631 assert!(stored.content.contains("<REDACTED>"));
1632 let hits = store
1633 .episodic()
1634 .search("rust retrieval", 10, false)
1635 .unwrap();
1636 assert_eq!(hits.len(), 1);
1637 assert_eq!(hits[0].record.id, public_id);
1638 }
1639
1640 #[test]
1641 fn typed_keys_retrieve_vocabulary_mismatch() {
1642 let tmp = tempdir().unwrap();
1643 let store = MemoryStore::open_default(tmp.path()).unwrap();
1644 let dog = sample_record(1, "My Golden Retriever loves the park");
1645 let other = sample_record(2, "I bought a yellow dress");
1646 let dog_id = dog.id;
1647 store.episodic().append_batch(&[dog, other]).unwrap();
1648 let hits = store
1649 .episodic()
1650 .search("What breed is my dog?", 5, false)
1651 .unwrap();
1652 assert_eq!(hits[0].record.id, dog_id);
1653 }
1654
1655 #[test]
1656 fn planner_boosts_temporal_date_match() {
1657 let tmp = tempdir().unwrap();
1658 let store = MemoryStore::open_default(tmp.path()).unwrap();
1659 let dated = sample_record(1, "I volunteered on February 14th at the animal shelter");
1660 let other = sample_record(2, "I volunteered at the community garden last summer");
1661 let dated_id = dated.id;
1662 store.episodic().append_batch(&[dated, other]).unwrap();
1663 let hits = store
1664 .episodic()
1665 .search("When did I volunteer at the animal shelter?", 5, false)
1666 .unwrap();
1667 assert_eq!(hits[0].record.id, dated_id);
1668 }
1669
1670 #[test]
1671 #[ignore = "manual in-process latency profile"]
1672 fn profile_ingest_and_search_latency() {
1673 fn timed_ms(label: &str, repeats: u32, mut work: impl FnMut()) {
1674 let start = std::time::Instant::now();
1675 for _ in 0..repeats {
1676 work();
1677 }
1678 let elapsed = start.elapsed();
1679 println!(
1680 "{label}: {:.3} ms (n={repeats})",
1681 elapsed.as_secs_f64() * 1000.0 / f64::from(repeats)
1682 );
1683 }
1684
1685 let search_records: Vec<EpisodicRecord> = (0..10_000)
1686 .map(|n| {
1687 sample_record(
1688 n,
1689 if n % 5 == 0 {
1690 "Rust memory retrieval benchmark item"
1691 } else {
1692 "Unrelated episodic record"
1693 },
1694 )
1695 })
1696 .collect();
1697
1698 timed_ms("append_single_1000", 1, || {
1699 let tmp = tempdir().unwrap();
1700 let store = MemoryStore::open_default(tmp.path()).unwrap();
1701 for n in 0..1_000 {
1702 store.episodic().append(&sample_record(n, "once")).unwrap();
1703 }
1704 });
1705 timed_ms("append_batch_1000", 1, || {
1706 let tmp = tempdir().unwrap();
1707 let store = MemoryStore::open_default(tmp.path()).unwrap();
1708 let records: Vec<EpisodicRecord> =
1709 (0..1_000).map(|n| sample_record(n, "once")).collect();
1710 store.episodic().append_batch(&records).unwrap();
1711 });
1712
1713 let tmp = tempdir().unwrap();
1714 {
1715 let store = MemoryStore::open_default(tmp.path()).unwrap();
1716 store.episodic().append_batch(&search_records).unwrap();
1717 }
1718 let cold = MemoryStore::open_default(tmp.path()).unwrap();
1719 timed_ms("cold_search_10000", 1, || {
1720 let hits = cold
1721 .episodic()
1722 .search("rust memory retrieval", 10, false)
1723 .unwrap();
1724 assert!(!hits.is_empty());
1725 });
1726 timed_ms("warm_search_10000", 50, || {
1727 let hits = cold
1728 .episodic()
1729 .search("rust memory retrieval", 10, false)
1730 .unwrap();
1731 assert!(!hits.is_empty());
1732 });
1733 }
1734
1735 #[test]
1741 #[ignore = "manual in-process latency profile at realistic scale"]
1742 fn profile_search_latency_25k_realistic() {
1743 const TOTAL: u64 = 25_000;
1744 const SESSIONS: u64 = 50;
1745 let topics = [
1746 "bookshelf",
1747 "guitar",
1748 "vegetarian",
1749 "portfolio",
1750 "commute",
1751 "grandmother",
1752 "chemistry",
1753 "marathon",
1754 "internship",
1755 "yoga",
1756 "spam filter",
1757 "projector",
1758 "swimming",
1759 "cousin",
1760 "bank account",
1761 "book club",
1762 "recipe",
1763 "journal subscription",
1764 "laptop",
1765 "hiking",
1766 ];
1767 let fillers = [
1768 "We discussed the plan for the weekend and agreed on the schedule.",
1769 "The meeting notes were circulated and everyone acknowledged them.",
1770 "I explained my reasoning and the group considered the proposal.",
1771 "After the presentation we reviewed the feedback together.",
1772 "She mentioned the deadline and we adjusted the timeline accordingly.",
1773 ];
1774
1775 let records: Vec<EpisodicRecord> = (0..TOTAL)
1776 .map(|n| {
1777 let session = n * SESSIONS / TOTAL;
1778 let topic = topics[(n as usize) % topics.len()];
1779 let filler = fillers[(n as usize) % fillers.len()];
1780 let content = format!(
1781 "Session {session} note {n}: my friend Alice mentioned {topic} while {filler}"
1782 );
1783 let session_id = if n % 4 == 0 {
1784 None
1785 } else {
1786 Some(uuid::Uuid::new_v4())
1787 };
1788 EpisodicRecord::new(
1789 session_id,
1790 n,
1791 if n % 3 == 0 {
1792 EpisodicKind::UserStatement
1793 } else {
1794 EpisodicKind::AssistantResponse
1795 },
1796 content,
1797 Provenance::new(ProvenanceSource::User),
1798 )
1799 })
1800 .collect();
1801
1802 let tmp = tempdir().unwrap();
1803 {
1804 let store = MemoryStore::open_default(tmp.path()).unwrap();
1805 let t = std::time::Instant::now();
1806 store.episodic().append_batch(&records).unwrap();
1807 println!("ingest 25k: {:.1} ms", t.elapsed().as_secs_f64() * 1000.0);
1808 }
1809
1810 let cold = MemoryStore::open_default(tmp.path()).unwrap();
1812 let episodic = cold.episodic();
1813 let t = std::time::Instant::now();
1814 let hits = episodic
1815 .search("guitar grandmother recipe", 10, false)
1816 .unwrap();
1817 println!(
1818 "cold_search: {:.1} ms (hits={})",
1819 t.elapsed().as_secs_f64() * 1000.0,
1820 hits.len()
1821 );
1822
1823 let queries: Vec<String> = (0..20)
1826 .map(|i| {
1827 let a = topics[i * 7 % topics.len()];
1828 let b = topics[(i * 7 + 5) % topics.len()];
1829 format!("{a} {b} weekend plan")
1830 })
1831 .collect();
1832 let start = std::time::Instant::now();
1833 for q in &queries {
1834 let hits = episodic.search(q, 10, false).unwrap();
1835 assert!(!hits.is_empty(), "no hits for {q}");
1836 }
1837 println!(
1838 "warm_search varied p50: {:.2} ms/query",
1839 start.elapsed().as_secs_f64() * 1000.0 / queries.len() as f64
1840 );
1841
1842 let t = std::time::Instant::now();
1843 let hits = episodic
1844 .search("zzzterm zzzother zzzthird", 10, false)
1845 .unwrap();
1846 println!(
1847 "no_match_query (full-scan fallback path): {:.1} ms (hits={})",
1848 t.elapsed().as_secs_f64() * 1000.0,
1849 hits.len()
1850 );
1851 }
1852
1853 #[test]
1854 fn enrichment_bridges_vocabulary_gap_for_theater() {
1855 let tmp = tempdir().unwrap();
1856 let store = MemoryStore::open_default(tmp.path()).unwrap();
1857 store.set_episodic_enrichment(crate::enrichment::VocabularyEnrichment::with_defaults());
1858 let answer = sample_record(1, "The production I attended was The Glass Menagerie");
1860 let competing = sample_record(2, "I went to a play at the local community theater");
1861 let answer_id = answer.id;
1862 store.episodic().append_batch(&[answer, competing]).unwrap();
1863 let hits = store
1864 .episodic()
1865 .search(
1866 "What play did I attend at the local community theater?",
1867 5,
1868 false,
1869 )
1870 .unwrap();
1871 assert!(hits.iter().any(|h| h.record.id == answer_id));
1874 }
1875
1876 #[test]
1877 fn enrichment_bridges_vocabulary_gap_for_shelter() {
1878 let tmp = tempdir().unwrap();
1879 let store = MemoryStore::open_default(tmp.path()).unwrap();
1880 store.set_episodic_enrichment(crate::enrichment::VocabularyEnrichment::with_defaults());
1881 let answer = sample_record(1, "I rescued a dog from the humane society last week");
1882 let other = sample_record(2, "I bought groceries at the store");
1883 let answer_id = answer.id;
1884 store.episodic().append_batch(&[answer, other]).unwrap();
1885 let hits = store
1886 .episodic()
1887 .search("When did I volunteer at the animal shelter?", 5, false)
1888 .unwrap();
1889 assert!(hits.iter().any(|h| h.record.id == answer_id));
1890 }
1891
1892 #[test]
1893 fn session_boost_favors_sessions_with_multiple_matches() {
1894 let tmp = tempdir().unwrap();
1895 let store = MemoryStore::open_default(tmp.path()).unwrap();
1896 let session_a = uuid::Uuid::new_v4();
1897 let session_b = uuid::Uuid::new_v4();
1898 let a1 = EpisodicRecord::new(
1900 Some(session_a),
1901 1,
1902 EpisodicKind::Observation,
1903 "I love hiking in the mountains",
1904 Provenance::new(ProvenanceSource::User),
1905 );
1906 let a2 = EpisodicRecord::new(
1907 Some(session_a),
1908 2,
1909 EpisodicKind::Observation,
1910 "Hiking in the mountains is great exercise",
1911 Provenance::new(ProvenanceSource::User),
1912 );
1913 let b1 = EpisodicRecord::new(
1915 Some(session_b),
1916 1,
1917 EpisodicKind::Observation,
1918 "Hiking is fun",
1919 Provenance::new(ProvenanceSource::User),
1920 );
1921 store.episodic().append_batch(&[a1, a2, b1]).unwrap();
1922 let hits = store
1923 .episodic()
1924 .search("hiking mountains", 10, false)
1925 .unwrap();
1926 let a_ranks: Vec<usize> = hits
1929 .iter()
1930 .enumerate()
1931 .filter(|(_, h)| h.record.session_id == Some(session_a))
1932 .map(|(i, _)| i)
1933 .collect();
1934 let b_rank = hits
1935 .iter()
1936 .position(|h| h.record.session_id == Some(session_b));
1937 if let Some(br) = b_rank {
1938 assert!(a_ranks.iter().all(|&ar| ar < br));
1939 }
1940 }
1941}