1use std::future::Future;
5use std::pin::Pin;
6use std::sync::Arc;
7use std::sync::atomic::{AtomicU64, Ordering};
8
9use futures::{StreamExt as _, TryStreamExt as _};
10use zeph_llm::provider::{LlmProvider as _, Message};
11
12const CHARS_PER_TOKEN: usize = 4;
14
15const CHUNK_CHARS: usize = 400 * CHARS_PER_TOKEN;
17
18const CHUNK_OVERLAP_CHARS: usize = 80 * CHARS_PER_TOKEN;
20
21fn chunk_text(text: &str) -> Vec<&str> {
27 if text.len() <= CHUNK_CHARS {
28 return vec![text];
29 }
30
31 let mut chunks = Vec::new();
32 let mut start = 0;
33
34 while start < text.len() {
35 let end = if start + CHUNK_CHARS >= text.len() {
36 text.len()
37 } else {
38 let boundary = text.floor_char_boundary(start + CHUNK_CHARS);
40 let slice = &text[start..boundary];
42 if let Some(pos) = slice.rfind("\n\n") {
43 start + pos + 2
44 } else if let Some(pos) = slice.rfind('\n') {
45 start + pos + 1
46 } else if let Some(pos) = slice.rfind(' ') {
47 start + pos + 1
48 } else {
49 boundary
50 }
51 };
52
53 chunks.push(&text[start..end]);
54 if end >= text.len() {
55 break;
56 }
57 let next = end.saturating_sub(CHUNK_OVERLAP_CHARS);
61 let new_start = text.ceil_char_boundary(next);
62 start = if new_start > start { new_start } else { end };
63 }
64
65 chunks
66}
67
68use crate::admission::{AdmissionDecision, log_admission_decision};
69use crate::embedding_store::{MessageKind, SearchFilter};
70use crate::error::MemoryError;
71use crate::store::admission_training::AdmissionTrainingInput;
72use crate::types::{ConversationId, MessageId};
73
74use super::SemanticMemory;
75use super::algorithms::{apply_mmr, apply_temporal_decay};
76
77#[derive(Debug, Clone, Default)]
81pub struct EmbedContext {
82 pub tool_name: Option<String>,
83 pub exit_code: Option<i32>,
84 pub timestamp: Option<String>,
85}
86
87#[derive(Debug)]
88pub struct RecalledMessage {
89 pub message: Message,
90 pub score: f32,
91}
92
93const MAX_EMBED_BG_TASKS: usize = 64;
95
96const QDRANT_WARN_WINDOW_SECS: u64 = 10;
98
99fn should_emit_qdrant_warn(last: u64, now: u64, window_secs: u64) -> bool {
101 now.saturating_sub(last) >= window_secs
102}
103
104fn warn_qdrant_ensure_failure(last_warn: &AtomicU64, log_tag: &str, err: &MemoryError) {
107 let now = std::time::SystemTime::now()
108 .duration_since(std::time::UNIX_EPOCH)
109 .unwrap_or_default()
110 .as_secs();
111 let last = last_warn.load(Ordering::Relaxed);
112 if should_emit_qdrant_warn(last, now, QDRANT_WARN_WINDOW_SECS) {
113 last_warn.store(now, Ordering::Relaxed);
114 tracing::warn!("{log_tag}: failed to ensure Qdrant collection: {err:#}");
115 } else {
116 tracing::debug!("{log_tag}: failed to ensure Qdrant collection (suppressed): {err:#}");
117 }
118}
119
120struct EmbedBgArgs {
126 qdrant: Arc<crate::embedding_store::EmbeddingStore>,
127 embed_provider: zeph_llm::any::AnyProvider,
128 message_id: MessageId,
129 content: String,
130 last_qdrant_warn: Arc<AtomicU64>,
131}
132
133async fn embed_chunk_and_store_bg<F>(args: EmbedBgArgs, log_tag: &'static str, store_chunk: F)
141where
142 F: Fn(u32, Vec<f32>) -> Pin<Box<dyn Future<Output = Result<(), MemoryError>> + Send>> + Send,
143{
144 let EmbedBgArgs {
145 qdrant,
146 embed_provider,
147 message_id,
148 content,
149 last_qdrant_warn,
150 } = args;
151 let chunks = chunk_text(&content);
152 let chunk_count = chunks.len();
153
154 let vectors = match embed_provider.embed_batch(&chunks).await {
155 Ok(v) => v,
156 Err(e) => {
157 tracing::warn!("{log_tag}: failed to embed chunks for msg {message_id}: {e:#}");
158 return;
159 }
160 };
161
162 let Some(first) = vectors.first() else {
163 return;
164 };
165 if let Err(e) = qdrant.ensure_collection_for_vector(first).await {
166 warn_qdrant_ensure_failure(&last_qdrant_warn, log_tag, &e);
167 return;
168 }
169
170 for (chunk_index, vector) in vectors.into_iter().enumerate() {
171 let chunk_index_u32 = u32::try_from(chunk_index).unwrap_or(u32::MAX);
172 if let Err(e) = store_chunk(chunk_index_u32, vector).await {
173 tracing::warn!(
174 "{log_tag}: failed to store chunk {chunk_index}/{chunk_count} \
175 for msg {message_id}: {e:#}"
176 );
177 }
178 }
179}
180
181enum AdmissionOutcome {
183 Reject,
185 Proceed(Option<AdmissionDecision>),
189}
190
191fn hela_outer_timeout(params: &crate::graph::HelaSpreadParams) -> std::time::Duration {
197 let embed_component = params
198 .embed_timeout
199 .unwrap_or(std::time::Duration::from_secs(5));
200 let step_stages = params.spread_depth.clamp(1, 6) + 2;
201 let step_component = params
202 .step_budget
203 .unwrap_or(std::time::Duration::from_millis(80))
204 * step_stages;
205 embed_component + step_component + std::time::Duration::from_millis(250)
206}
207
208impl SemanticMemory {
209 #[cfg_attr(
219 feature = "profiling",
220 tracing::instrument(name = "memory.remember", skip_all, fields(content_len = %content.len()))
221 )]
222 pub async fn remember(
223 &self,
224 conversation_id: ConversationId,
225 role: &str,
226 content: &str,
227 goal_text: Option<&str>,
228 ) -> Result<Option<MessageId>, MemoryError> {
229 let admission_decision = match self
230 .run_admission_gate(conversation_id, role, content, goal_text)
231 .await
232 {
233 AdmissionOutcome::Reject => return Ok(None),
234 AdmissionOutcome::Proceed(decision) => decision,
235 };
236
237 if self
238 .run_quality_gate(conversation_id, role, content, admission_decision.as_ref())
239 .await
240 {
241 return Ok(None);
242 }
243
244 let message_id = self
245 .sqlite
246 .save_message(conversation_id, role, content)
247 .await?;
248
249 self.record_admission_sample_opt(
250 conversation_id,
251 role,
252 content,
253 admission_decision.as_ref(),
254 Some(message_id),
255 )
256 .await;
257
258 self.embed_and_store_regular(message_id, conversation_id, role, content);
259
260 Ok(Some(message_id))
261 }
262
263 #[cfg_attr(
272 feature = "profiling",
273 tracing::instrument(name = "memory.remember", skip_all, fields(content_len = %content.len()))
274 )]
275 pub async fn remember_with_parts(
276 &self,
277 conversation_id: ConversationId,
278 role: &str,
279 content: &str,
280 parts_json: &str,
281 goal_text: Option<&str>,
282 ) -> Result<(Option<MessageId>, bool), MemoryError> {
283 let admission_decision = match self
284 .run_admission_gate(conversation_id, role, content, goal_text)
285 .await
286 {
287 AdmissionOutcome::Reject => return Ok((None, false)),
288 AdmissionOutcome::Proceed(decision) => decision,
289 };
290
291 if self
292 .run_quality_gate(conversation_id, role, content, admission_decision.as_ref())
293 .await
294 {
295 return Ok((None, false));
296 }
297
298 let message_id = self
299 .sqlite
300 .save_message_with_parts(conversation_id, role, content, parts_json)
301 .await?;
302
303 self.record_admission_sample_opt(
304 conversation_id,
305 role,
306 content,
307 admission_decision.as_ref(),
308 Some(message_id),
309 )
310 .await;
311
312 let embedding_stored =
313 self.embed_and_store_regular(message_id, conversation_id, role, content);
314
315 Ok((Some(message_id), embedding_stored))
316 }
317
318 #[cfg_attr(
330 feature = "profiling",
331 tracing::instrument(name = "memory.remember", skip_all, fields(content_len = %content.len()))
332 )]
333 pub async fn remember_tool_output(
334 &self,
335 conversation_id: ConversationId,
336 role: &str,
337 content: &str,
338 parts_json: &str,
339 embed_ctx: EmbedContext,
340 ) -> Result<(Option<MessageId>, bool), MemoryError> {
341 let admission_decision = match self
344 .run_admission_gate(conversation_id, role, content, None)
345 .await
346 {
347 AdmissionOutcome::Reject => return Ok((None, false)),
348 AdmissionOutcome::Proceed(decision) => decision,
349 };
350
351 let message_id = self
352 .sqlite
353 .save_message_with_parts(conversation_id, role, content, parts_json)
354 .await?;
355
356 self.record_admission_sample_opt(
357 conversation_id,
358 role,
359 content,
360 admission_decision.as_ref(),
361 Some(message_id),
362 )
363 .await;
364
365 let embedding_stored = self.embed_chunks_with_tool_context(
366 message_id,
367 conversation_id,
368 role,
369 content,
370 embed_ctx,
371 );
372
373 Ok((Some(message_id), embedding_stored))
374 }
375
376 #[cfg_attr(
387 feature = "profiling",
388 tracing::instrument(name = "memory.remember", skip_all, fields(content_len = %content.len()))
389 )]
390 pub async fn remember_categorized(
391 &self,
392 conversation_id: ConversationId,
393 role: &str,
394 content: &str,
395 category: Option<&str>,
396 goal_text: Option<&str>,
397 ) -> Result<Option<MessageId>, MemoryError> {
398 let admission_decision = match self
401 .run_admission_gate(conversation_id, role, content, goal_text)
402 .await
403 {
404 AdmissionOutcome::Reject => return Ok(None),
405 AdmissionOutcome::Proceed(decision) => decision,
406 };
407
408 let message_id = self
409 .sqlite
410 .save_message_with_category(conversation_id, role, content, category)
411 .await?;
412
413 self.record_admission_sample_opt(
414 conversation_id,
415 role,
416 content,
417 admission_decision.as_ref(),
418 Some(message_id),
419 )
420 .await;
421
422 self.embed_and_store_with_category(message_id, conversation_id, role, content, category);
423
424 Ok(Some(message_id))
425 }
426
427 async fn run_admission_gate(
433 &self,
434 conversation_id: ConversationId,
435 role: &str,
436 content: &str,
437 goal_text: Option<&str>,
438 ) -> AdmissionOutcome {
439 let Some(admission) = &self.admission_control else {
440 return AdmissionOutcome::Proceed(None);
441 };
442 let decision = admission
443 .evaluate(
444 content,
445 role,
446 self.effective_embed_provider(),
447 self.qdrant.as_ref(),
448 goal_text,
449 )
450 .await;
451 let preview: String = content.chars().take(100).collect();
452 log_admission_decision(&decision, &preview, role, admission.threshold());
453 if !decision.admitted {
454 self.record_admission_sample(conversation_id, role, content, &decision, None)
455 .await;
456 return AdmissionOutcome::Reject;
457 }
458 AdmissionOutcome::Proceed(Some(decision))
459 }
460
461 async fn run_quality_gate(
468 &self,
469 conversation_id: ConversationId,
470 role: &str,
471 content: &str,
472 admission_decision: Option<&AdmissionDecision>,
473 ) -> bool {
474 let Some(gate) = &self.quality_gate else {
475 return false;
476 };
477 if gate
478 .evaluate(content, self.effective_embed_provider(), &[])
479 .await
480 .is_none()
481 {
482 return false;
483 }
484 if let Some(decision) = admission_decision {
485 self.record_admission_sample(conversation_id, role, content, decision, None)
486 .await;
487 }
488 true
489 }
490
491 async fn record_admission_sample_opt(
494 &self,
495 conversation_id: ConversationId,
496 role: &str,
497 content: &str,
498 decision: Option<&AdmissionDecision>,
499 message_id: Option<MessageId>,
500 ) {
501 if let Some(decision) = decision {
502 self.record_admission_sample(conversation_id, role, content, decision, message_id)
503 .await;
504 }
505 }
506
507 async fn record_admission_sample(
515 &self,
516 conversation_id: ConversationId,
517 role: &str,
518 content: &str,
519 decision: &AdmissionDecision,
520 message_id: Option<MessageId>,
521 ) {
522 let features_json = match serde_json::to_string(&decision.factors) {
523 Ok(json) => json,
524 Err(e) => {
525 tracing::debug!(error = %e, "admission training: failed to serialize factors");
526 return;
527 }
528 };
529 if let Err(e) = self
530 .sqlite
531 .record_admission_training(AdmissionTrainingInput {
532 message_id,
533 conversation_id,
534 content,
535 role,
536 composite_score: decision.composite_score,
537 was_admitted: decision.admitted,
538 features_json: &features_json,
539 })
540 .await
541 {
542 tracing::debug!(error = %e, "admission training: failed to record sample (non-fatal)");
543 }
544 }
545
546 pub async fn recall_with_category(
554 &self,
555 query: &str,
556 limit: usize,
557 filter: Option<SearchFilter>,
558 category: Option<&str>,
559 ) -> Result<Vec<RecalledMessage>, MemoryError> {
560 let filter_with_category = filter.map(|mut f| {
561 f.category = category.map(str::to_owned);
562 f
563 });
564 self.recall(query, limit, filter_with_category).await
565 }
566
567 pub fn reap_embed_tasks(&self) {
571 if let Ok(mut tasks) = self.embed_tasks.lock() {
572 while tasks.try_join_next().is_some() {}
573 }
574 }
575
576 fn spawn_embed_bg<F>(&self, fut: F) -> bool
580 where
581 F: std::future::Future<Output = ()> + Send + 'static,
582 {
583 let Ok(mut tasks) = self.embed_tasks.lock() else {
584 return false;
585 };
586 while tasks.try_join_next().is_some() {}
588 if tasks.len() >= MAX_EMBED_BG_TASKS {
589 tracing::debug!("background embed task limit reached, skipping");
590 return false;
591 }
592 tasks.spawn(fut);
593 true
594 }
595
596 fn embed_and_store_with_category(
600 &self,
601 message_id: MessageId,
602 conversation_id: ConversationId,
603 role: &str,
604 content: &str,
605 category: Option<&str>,
606 ) -> bool {
607 let Some(qdrant) = self.qdrant.clone() else {
608 return false;
609 };
610 let embed_provider = self.effective_embed_provider().clone();
611 if !embed_provider.supports_embeddings() {
612 return false;
613 }
614 let store_qdrant = Arc::clone(&qdrant);
615 let embedding_model = self.embedding_model.clone();
616 let role = role.to_owned();
617 let category = category.map(str::to_owned);
618 let store_chunk =
619 move |chunk_index: u32,
620 vector: Vec<f32>|
621 -> Pin<Box<dyn Future<Output = Result<(), MemoryError>> + Send>> {
622 let qdrant = Arc::clone(&store_qdrant);
623 let embedding_model = embedding_model.clone();
624 let role = role.clone();
625 let category = category.clone();
626 Box::pin(async move {
627 qdrant
628 .store_with_category(
629 message_id,
630 conversation_id,
631 &role,
632 vector,
633 MessageKind::Regular,
634 &embedding_model,
635 chunk_index,
636 category.as_deref(),
637 )
638 .await
639 .map(|_| ())
640 })
641 };
642 self.spawn_embed_bg(embed_chunk_and_store_bg(
643 EmbedBgArgs {
644 qdrant,
645 embed_provider,
646 message_id,
647 content: content.to_owned(),
648 last_qdrant_warn: Arc::clone(&self.last_qdrant_warn),
649 },
650 "bg embed_category",
651 store_chunk,
652 ))
653 }
654
655 fn embed_and_store_regular(
659 &self,
660 message_id: MessageId,
661 conversation_id: ConversationId,
662 role: &str,
663 content: &str,
664 ) -> bool {
665 let Some(qdrant) = self.qdrant.clone() else {
666 return false;
667 };
668 let embed_provider = self.effective_embed_provider().clone();
669 if !embed_provider.supports_embeddings() {
670 return false;
671 }
672 let store_qdrant = Arc::clone(&qdrant);
673 let embedding_model = self.embedding_model.clone();
674 let role = role.to_owned();
675 let store_chunk =
676 move |chunk_index: u32,
677 vector: Vec<f32>|
678 -> Pin<Box<dyn Future<Output = Result<(), MemoryError>> + Send>> {
679 let qdrant = Arc::clone(&store_qdrant);
680 let embedding_model = embedding_model.clone();
681 let role = role.clone();
682 Box::pin(async move {
683 qdrant
684 .store(
685 message_id,
686 conversation_id,
687 &role,
688 vector,
689 MessageKind::Regular,
690 &embedding_model,
691 chunk_index,
692 )
693 .await
694 .map(|_| ())
695 })
696 };
697 self.spawn_embed_bg(embed_chunk_and_store_bg(
698 EmbedBgArgs {
699 qdrant,
700 embed_provider,
701 message_id,
702 content: content.to_owned(),
703 last_qdrant_warn: Arc::clone(&self.last_qdrant_warn),
704 },
705 "bg embed_regular",
706 store_chunk,
707 ))
708 }
709
710 fn embed_chunks_with_tool_context(
714 &self,
715 message_id: MessageId,
716 conversation_id: ConversationId,
717 role: &str,
718 content: &str,
719 embed_ctx: EmbedContext,
720 ) -> bool {
721 let Some(qdrant) = self.qdrant.clone() else {
722 return false;
723 };
724 let embed_provider = self.effective_embed_provider().clone();
725 if !embed_provider.supports_embeddings() {
726 return false;
727 }
728 let store_qdrant = Arc::clone(&qdrant);
729 let embedding_model = self.embedding_model.clone();
730 let role = role.to_owned();
731 let store_chunk =
732 move |chunk_index: u32,
733 vector: Vec<f32>|
734 -> Pin<Box<dyn Future<Output = Result<(), MemoryError>> + Send>> {
735 let qdrant = Arc::clone(&store_qdrant);
736 let embedding_model = embedding_model.clone();
737 let role = role.clone();
738 let embed_ctx = embed_ctx.clone();
739 Box::pin(async move {
740 if let Some(tool_name) = embed_ctx.tool_name {
741 qdrant
742 .store_with_tool_context(
743 message_id,
744 conversation_id,
745 &role,
746 vector,
747 MessageKind::Regular,
748 &embedding_model,
749 chunk_index,
750 &tool_name,
751 embed_ctx.exit_code,
752 embed_ctx.timestamp.as_deref(),
753 )
754 .await
755 .map(|_| ())
756 } else {
757 qdrant
758 .store(
759 message_id,
760 conversation_id,
761 &role,
762 vector,
763 MessageKind::Regular,
764 &embedding_model,
765 chunk_index,
766 )
767 .await
768 .map(|_| ())
769 }
770 })
771 };
772 self.spawn_embed_bg(embed_chunk_and_store_bg(
773 EmbedBgArgs {
774 qdrant,
775 embed_provider,
776 message_id,
777 content: content.to_owned(),
778 last_qdrant_warn: Arc::clone(&self.last_qdrant_warn),
779 },
780 "bg embed_tool",
781 store_chunk,
782 ))
783 }
784
785 pub async fn save_only(
793 &self,
794 conversation_id: ConversationId,
795 role: &str,
796 content: &str,
797 parts_json: &str,
798 ) -> Result<MessageId, MemoryError> {
799 self.sqlite
800 .save_message_with_parts(conversation_id, role, content, parts_json)
801 .await
802 }
803
804 #[cfg_attr(
814 feature = "profiling",
815 tracing::instrument(name = "memory.recall", skip_all, fields(query_len = %query.len(), result_count = tracing::field::Empty, top_score = tracing::field::Empty))
816 )]
817 pub async fn recall(
818 &self,
819 query: &str,
820 limit: usize,
821 filter: Option<SearchFilter>,
822 ) -> Result<Vec<RecalledMessage>, MemoryError> {
823 let conversation_id = filter.as_ref().and_then(|f| f.conversation_id);
824
825 tracing::debug!(
826 query_len = query.len(),
827 limit,
828 has_filter = filter.is_some(),
829 conversation_id = conversation_id.map(|c| c.0),
830 has_qdrant = self.qdrant.is_some(),
831 "recall: starting hybrid search"
832 );
833
834 let keyword_results = match self
835 .sqlite
836 .keyword_search(query, self.effective_depth(limit), conversation_id)
837 .await
838 {
839 Ok(results) => results,
840 Err(e) => {
841 tracing::warn!("FTS5 keyword search failed: {e:#}");
842 Vec::new()
843 }
844 };
845
846 let vector_results = if let Some(qdrant) = &self.qdrant
847 && self.effective_embed_provider().supports_embeddings()
848 {
849 let embed_input = self.apply_search_prompt(query);
850 let query_vector = match tokio::time::timeout(
851 self.embed_timeout,
852 self.effective_embed_provider().embed(&embed_input),
853 )
854 .await
855 {
856 Ok(Ok(v)) => v,
857 Ok(Err(e)) => return Err(e.into()),
858 Err(_) => {
859 tracing::warn!("recall_semantic: embed timed out, returning empty results");
860 return Ok(Vec::new());
861 }
862 };
863 let query_vector = self.apply_query_bias(query, query_vector).await;
864 qdrant.ensure_collection_for_vector(&query_vector).await?;
865 qdrant
866 .search(&query_vector, self.effective_depth(limit), filter)
867 .await?
868 } else {
869 Vec::new()
870 };
871
872 let results = self
873 .recall_merge_and_rank(keyword_results, vector_results, limit, None)
874 .await?;
875 #[cfg(feature = "profiling")]
876 {
877 let span = tracing::Span::current();
878 span.record("result_count", results.len());
879 if let Some(top) = results.first() {
880 span.record("top_score", top.score);
881 }
882 }
883 Ok(results)
884 }
885
886 #[cfg_attr(
887 feature = "profiling",
888 tracing::instrument(name = "memory.recall.fts5", skip_all, fields(query_len = %query.len()))
889 )]
890 pub(super) async fn recall_fts5_raw(
891 &self,
892 query: &str,
893 limit: usize,
894 conversation_id: Option<ConversationId>,
895 ) -> Result<Vec<(MessageId, f64)>, MemoryError> {
896 self.sqlite
897 .keyword_search(query, self.effective_depth(limit), conversation_id)
898 .await
899 }
900
901 #[cfg_attr(
902 feature = "profiling",
903 tracing::instrument(name = "memory.recall.vectors", skip_all, fields(query_len = %query.len()))
904 )]
905 pub(super) async fn recall_vectors_raw(
906 &self,
907 query: &str,
908 limit: usize,
909 filter: Option<SearchFilter>,
910 ) -> Result<Vec<crate::embedding_store::SearchResult>, MemoryError> {
911 let Some(qdrant) = &self.qdrant else {
912 return Ok(Vec::new());
913 };
914 if !self.effective_embed_provider().supports_embeddings() {
915 return Ok(Vec::new());
916 }
917 let embed_input = self.apply_search_prompt(query);
918 let query_vector = match tokio::time::timeout(
919 self.embed_timeout,
920 self.effective_embed_provider().embed(&embed_input),
921 )
922 .await
923 {
924 Ok(Ok(v)) => v,
925 Ok(Err(e)) => return Err(e.into()),
926 Err(_) => {
927 tracing::warn!("recall_vectors_raw: embed timed out, returning empty results");
928 return Ok(Vec::new());
929 }
930 };
931 let query_vector = self.apply_query_bias(query, query_vector).await;
932 qdrant.ensure_collection_for_vector(&query_vector).await?;
933 qdrant
934 .search(&query_vector, self.effective_depth(limit), filter)
935 .await
936 }
937
938 #[cfg_attr(
947 feature = "profiling",
948 tracing::instrument(name = "memory.recall.merge_and_rank", skip_all, fields(kw_count = keyword_results.len(), vec_count = vector_results.len()))
949 )]
950 #[allow(clippy::cast_possible_truncation, clippy::too_many_lines)]
951 pub(super) async fn recall_merge_and_rank(
952 &self,
953 keyword_results: Vec<(MessageId, f64)>,
954 vector_results: Vec<crate::embedding_store::SearchResult>,
955 limit: usize,
956 goal_entity_id: Option<i64>,
957 ) -> Result<Vec<RecalledMessage>, MemoryError> {
958 tracing::debug!(
959 vector_count = vector_results.len(),
960 keyword_count = keyword_results.len(),
961 limit,
962 "recall: merging search results"
963 );
964
965 let mut scores: std::collections::HashMap<MessageId, f64> =
966 std::collections::HashMap::new();
967
968 if !vector_results.is_empty() {
969 let max_vs = vector_results
970 .iter()
971 .map(|r| r.score)
972 .fold(f32::NEG_INFINITY, f32::max);
973 let norm = if max_vs > 0.0 { max_vs } else { 1.0 };
974 for r in &vector_results {
975 let normalized = f64::from(r.score / norm);
976 *scores.entry(r.message_id).or_default() += normalized * self.vector_weight;
977 }
978 }
979
980 if !keyword_results.is_empty() {
981 let max_ks = keyword_results
982 .iter()
983 .map(|r| r.1)
984 .fold(f64::NEG_INFINITY, f64::max);
985 let norm = if max_ks > 0.0 { max_ks } else { 1.0 };
986 for &(msg_id, score) in &keyword_results {
987 let normalized = score / norm;
988 *scores.entry(msg_id).or_default() += normalized * self.keyword_weight;
989 }
990 }
991
992 if scores.is_empty() {
993 tracing::debug!("recall: empty merge, no overlapping scores");
994 return Ok(Vec::new());
995 }
996
997 let mut ranked: Vec<(MessageId, f64)> = scores.into_iter().collect();
998 ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
999
1000 tracing::debug!(
1001 merged = ranked.len(),
1002 top_score = ranked.first().map(|r| r.1),
1003 bottom_score = ranked.last().map(|r| r.1),
1004 vector_weight = %self.vector_weight,
1005 keyword_weight = %self.keyword_weight,
1006 "recall: weighted merge complete"
1007 );
1008
1009 if self.temporal_decay.is_enabled() && self.temporal_decay_half_life_days > 0 {
1010 let ids: Vec<MessageId> = ranked.iter().map(|r| r.0).collect();
1011 match self.sqlite.message_timestamps(&ids).await {
1012 Ok(timestamps) => {
1013 apply_temporal_decay(
1014 &mut ranked,
1015 ×tamps,
1016 self.temporal_decay_half_life_days,
1017 );
1018 ranked
1019 .sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
1020 tracing::debug!(
1021 half_life_days = self.temporal_decay_half_life_days,
1022 top_score_after = ranked.first().map(|r| r.1),
1023 "recall: temporal decay applied"
1024 );
1025 }
1026 Err(e) => {
1027 tracing::warn!("temporal decay: failed to fetch timestamps: {e:#}");
1028 }
1029 }
1030 }
1031
1032 if self.mmr_reranking.is_enabled() && !vector_results.is_empty() {
1033 if let Some(qdrant) = &self.qdrant {
1034 let ids: Vec<MessageId> = ranked.iter().map(|r| r.0).collect();
1035 match qdrant.get_vectors(&ids).await {
1036 Ok(vec_map) if !vec_map.is_empty() => {
1037 let ranked_len_before = ranked.len();
1038 ranked = apply_mmr(&ranked, &vec_map, self.mmr_lambda, limit);
1039 tracing::debug!(
1040 before = ranked_len_before,
1041 after = ranked.len(),
1042 lambda = %self.mmr_lambda,
1043 "recall: mmr re-ranked"
1044 );
1045 }
1046 Ok(_) => {
1047 ranked.truncate(limit);
1048 }
1049 Err(e) => {
1050 tracing::warn!("MMR: failed to fetch vectors: {e:#}");
1051 ranked.truncate(limit);
1052 }
1053 }
1054 } else {
1055 ranked.truncate(limit);
1056 }
1057 } else {
1058 ranked.truncate(limit);
1059 }
1060
1061 if self.importance_scoring.is_enabled() && !ranked.is_empty() {
1062 let ids: Vec<MessageId> = ranked.iter().map(|r| r.0).collect();
1063 match self.sqlite.fetch_importance_scores(&ids).await {
1064 Ok(scores) => {
1065 for (msg_id, score) in &mut ranked {
1066 if let Some(&imp) = scores.get(msg_id) {
1067 *score += imp * self.importance_weight;
1068 }
1069 }
1070 ranked
1071 .sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
1072 tracing::debug!(
1073 importance_weight = %self.importance_weight,
1074 "recall: importance scores blended"
1075 );
1076 }
1077 Err(e) => {
1078 tracing::warn!("importance scoring: failed to fetch scores: {e:#}");
1079 }
1080 }
1081 }
1082
1083 if (self.tier_boost_semantic - 1.0).abs() > f64::EPSILON && !ranked.is_empty() {
1087 let ids: Vec<MessageId> = ranked.iter().map(|r| r.0).collect();
1088 match self.sqlite.fetch_tiers(&ids).await {
1089 Ok(tiers) => {
1090 let bonus = self.tier_boost_semantic - 1.0;
1091 let mut boosted = false;
1092 for (msg_id, score) in &mut ranked {
1093 if tiers.get(msg_id).map(String::as_str) == Some("semantic") {
1094 *score += bonus;
1095 boosted = true;
1096 }
1097 }
1098 if boosted {
1099 ranked.sort_by(|a, b| {
1100 b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)
1101 });
1102 tracing::debug!(
1103 tier_boost = %self.tier_boost_semantic,
1104 "recall: semantic tier boost applied"
1105 );
1106 }
1107 }
1108 Err(e) => {
1109 tracing::warn!("tier boost: failed to fetch tiers: {e:#}");
1110 }
1111 }
1112 }
1113
1114 if let Some(fs) = &self.five_signal
1116 && !fs.weights.is_baseline()
1117 {
1118 self.apply_five_signal_scoring(&mut ranked, fs, goal_entity_id)
1119 .await;
1120 }
1121
1122 let ids: Vec<MessageId> = ranked.iter().map(|r| r.0).collect();
1123
1124 if let Some(fs) = &self.five_signal {
1126 for id in &ids {
1127 fs.access_cache
1128 .log_access(*id, "message", &fs.session_id)
1129 .await;
1130 }
1131 fs.metrics.inc_recall();
1132 }
1133
1134 if !ids.is_empty()
1135 && let Err(e) = self.batch_increment_access_count(ids.clone()).await
1136 {
1137 tracing::warn!("recall: failed to increment access counts: {e:#}");
1138 }
1139
1140 if let Err(e) = self.sqlite.mark_training_recalled(&ids).await {
1142 tracing::debug!(
1143 error = %e,
1144 "recall: failed to mark training data as recalled (non-fatal)"
1145 );
1146 }
1147
1148 let messages = self.sqlite.messages_by_ids(&ids).await?;
1149 let msg_map: std::collections::HashMap<MessageId, _> = messages.into_iter().collect();
1150
1151 let recalled: Vec<RecalledMessage> = ranked
1152 .iter()
1153 .filter_map(|(msg_id, score)| {
1154 msg_map.get(msg_id).map(|msg| RecalledMessage {
1155 message: msg.clone(),
1156 #[expect(clippy::cast_possible_truncation)]
1157 score: *score as f32,
1158 })
1159 })
1160 .collect();
1161
1162 tracing::debug!(final_count = recalled.len(), "recall: final results");
1163
1164 Ok(recalled)
1165 }
1166
1167 async fn apply_five_signal_scoring(
1173 &self,
1174 ranked: &mut [(MessageId, f64)],
1175 fs: &crate::five_signal::FiveSignalRuntime,
1176 goal_entity_id: Option<i64>,
1177 ) {
1178 use crate::five_signal::causal_distance::CausalDistanceComputer;
1179 use crate::five_signal::scoring::{CandidateSignals, apply_five_signal_scoring};
1180 use sqlx::Row as _;
1181
1182 let ids: Vec<MessageId> = ranked.iter().map(|r| r.0).collect();
1183
1184 let freq_map = match fs
1186 .access_cache
1187 .load_for_candidates(&fs.session_id, &ids)
1188 .await
1189 {
1190 Ok(m) => m,
1191 Err(e) => {
1192 tracing::warn!(error = %e, "five_signal: failed to load access frequencies (skipping)");
1193 return;
1194 }
1195 };
1196
1197 let created_at_map: std::collections::HashMap<MessageId, i64> = {
1199 let id_vals: Vec<i64> = ids.iter().map(|id| id.0).collect();
1200 let placeholders = zeph_db::placeholder_list(1, id_vals.len());
1201 let created_at_epoch =
1202 <zeph_db::ActiveDialect as zeph_db::dialect::Dialect>::epoch_from_col("created_at");
1203 let sql = format!(
1204 "SELECT id, {created_at_epoch} AS created_at FROM messages \
1205 WHERE id IN ({placeholders}) AND deleted_at IS NULL"
1206 );
1207 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1208 for id in &id_vals {
1209 q = q.bind(id);
1210 }
1211 match q.fetch_all(&fs.pool).await {
1212 Ok(rows) => rows
1213 .iter()
1214 .map(|row| {
1215 (
1216 MessageId(row.get::<i64, _>("id")),
1217 row.get::<i64, _>("created_at"),
1218 )
1219 })
1220 .collect(),
1221 Err(e) => {
1222 tracing::warn!(error = %e, "five_signal: failed to fetch created_at (skipping novelty)");
1223 std::collections::HashMap::new()
1224 }
1225 }
1226 };
1227
1228 let causal_distance_map: std::collections::HashMap<i64, u32> = {
1232 let entity_ids: Vec<i64> = ids.iter().map(|id| id.0).collect();
1233 match fs
1234 .causal_computer
1235 .compute(goal_entity_id, &entity_ids)
1236 .await
1237 {
1238 Ok(m) => m,
1239 Err(e) => {
1240 tracing::warn!(error = %e, "five_signal: causal BFS failed (using neutral)");
1241 std::collections::HashMap::new()
1242 }
1243 }
1244 };
1245 let neutral_causal_score =
1246 CausalDistanceComputer::distance_to_score(fs.config.neutral_causal_distance);
1247
1248 let mut signals_map = std::collections::HashMap::with_capacity(ids.len());
1249 for &(msg_id, base_score) in ranked.iter() {
1250 let frequency = freq_map.get(&msg_id).copied().unwrap_or(0.0);
1251 let half = base_score / 2.0;
1254 let fact_created_at = created_at_map
1255 .get(&msg_id)
1256 .copied()
1257 .unwrap_or(fs.session_start);
1258 let novelty = fs.novelty_computer.compute(fact_created_at);
1259 let causal = causal_distance_map
1260 .get(&msg_id.0)
1261 .map_or(neutral_causal_score, |&d| {
1262 CausalDistanceComputer::distance_to_score(d)
1263 });
1264 signals_map.insert(
1265 msg_id,
1266 CandidateSignals {
1267 recency: half,
1268 relevance: half,
1269 frequency,
1270 causal,
1271 novelty,
1272 },
1273 );
1274 }
1275
1276 apply_five_signal_scoring(ranked, &fs.weights, &signals_map);
1277
1278 tracing::debug!(
1279 candidate_count = ids.len(),
1280 "recall: five-signal scoring applied"
1281 );
1282 }
1283
1284 async fn recall_by_route(
1290 &self,
1291 route: crate::router::MemoryRoute,
1292 query: &str,
1293 limit: usize,
1294 filter: Option<crate::embedding_store::SearchFilter>,
1295 ) -> Result<
1296 (
1297 Vec<(crate::types::MessageId, f64)>,
1298 Vec<crate::embedding_store::SearchResult>,
1299 ),
1300 MemoryError,
1301 > {
1302 use crate::router::MemoryRoute;
1303
1304 let conversation_id = filter.as_ref().and_then(|f| f.conversation_id);
1305
1306 let results: (
1307 Vec<(crate::types::MessageId, f64)>,
1308 Vec<crate::embedding_store::SearchResult>,
1309 ) = match route {
1310 MemoryRoute::Keyword => {
1311 let kw = self.recall_fts5_raw(query, limit, conversation_id).await?;
1312 (kw, Vec::new())
1313 }
1314 MemoryRoute::Hybrid => {
1315 let kw = match self.recall_fts5_raw(query, limit, conversation_id).await {
1316 Ok(r) => r,
1317 Err(e) => {
1318 tracing::warn!("FTS5 keyword search failed: {e:#}");
1319 Vec::new()
1320 }
1321 };
1322 let vr = self.recall_vectors_raw(query, limit, filter).await?;
1323 (kw, vr)
1324 }
1325 MemoryRoute::Episodic => {
1334 let range = crate::router::resolve_temporal_range(query, chrono::Utc::now());
1335 let cleaned = crate::router::strip_temporal_keywords(query);
1336 let search_query = if cleaned.is_empty() { query } else { &cleaned };
1337 let kw = if let Some(ref r) = range {
1338 self.sqlite
1339 .keyword_search_with_time_range(
1340 search_query,
1341 limit,
1342 conversation_id,
1343 r.after.as_deref(),
1344 r.before.as_deref(),
1345 )
1346 .await?
1347 } else {
1348 self.recall_fts5_raw(search_query, limit, conversation_id)
1349 .await?
1350 };
1351 tracing::debug!(
1352 has_range = range.is_some(),
1353 cleaned_query = %search_query,
1354 keyword_count = kw.len(),
1355 "recall: episodic path"
1356 );
1357 (kw, Vec::new())
1358 }
1359 MemoryRoute::Graph => {
1362 let kw = match self.recall_fts5_raw(query, limit, conversation_id).await {
1363 Ok(r) => r,
1364 Err(e) => {
1365 tracing::warn!("FTS5 keyword search failed (graph→hybrid fallback): {e:#}");
1366 Vec::new()
1367 }
1368 };
1369 let vr = self.recall_vectors_raw(query, limit, filter).await?;
1370 (kw, vr)
1371 }
1372 _ => {
1373 let vr = self.recall_vectors_raw(query, limit, filter).await?;
1374 (Vec::new(), vr)
1375 }
1376 };
1377 Ok(results)
1378 }
1379
1380 #[cfg_attr(
1392 feature = "profiling",
1393 tracing::instrument(name = "memory.recall", skip_all, fields(query_len = %query.len(), result_count = tracing::field::Empty))
1394 )]
1395 pub async fn recall_routed(
1396 &self,
1397 query: &str,
1398 limit: usize,
1399 filter: Option<SearchFilter>,
1400 router: &dyn crate::router::MemoryRouter,
1401 goal_entity_id: Option<i64>,
1402 ) -> Result<Vec<RecalledMessage>, MemoryError> {
1403 let route = router.route(query);
1404 tracing::debug!(?route, query_len = query.len(), "memory routing decision");
1405
1406 let (keyword_results, vector_results) =
1407 self.recall_by_route(route, query, limit, filter).await?;
1408
1409 tracing::debug!(
1410 keyword_count = keyword_results.len(),
1411 vector_count = vector_results.len(),
1412 "recall: routed search results"
1413 );
1414
1415 self.recall_merge_and_rank(keyword_results, vector_results, limit, goal_entity_id)
1416 .await
1417 }
1418
1419 #[cfg_attr(
1433 feature = "profiling",
1434 tracing::instrument(name = "memory.recall", skip_all, fields(query_len = %query.len(), result_count = tracing::field::Empty))
1435 )]
1436 pub async fn recall_routed_async(
1437 &self,
1438 query: &str,
1439 limit: usize,
1440 filter: Option<crate::embedding_store::SearchFilter>,
1441 router: &dyn crate::router::AsyncMemoryRouter,
1442 goal_entity_id: Option<i64>,
1443 ) -> Result<Vec<RecalledMessage>, MemoryError> {
1444 let decision = router.route_async(query).await;
1445 let route = decision.route;
1446 tracing::debug!(
1447 ?route,
1448 confidence = decision.confidence,
1449 query_len = query.len(),
1450 "memory routing decision (async)"
1451 );
1452
1453 let (keyword_results, vector_results) =
1454 self.recall_by_route(route, query, limit, filter).await?;
1455
1456 tracing::debug!(
1457 keyword_count = keyword_results.len(),
1458 vector_count = vector_results.len(),
1459 "recall: routed search results (async)"
1460 );
1461
1462 self.recall_merge_and_rank(keyword_results, vector_results, limit, goal_entity_id)
1463 .await
1464 }
1465
1466 #[cfg_attr(
1480 feature = "profiling",
1481 tracing::instrument(name = "memory.recall_graph", skip_all, fields(result_count = tracing::field::Empty))
1482 )]
1483 pub async fn recall_graph(
1484 &self,
1485 query: &str,
1486 limit: usize,
1487 max_hops: u32,
1488 at_timestamp: Option<&str>,
1489 temporal_decay_rate: f64,
1490 edge_types: &[crate::graph::EdgeType],
1491 ) -> Result<Vec<crate::graph::types::GraphFact>, MemoryError> {
1492 let Some(store) = &self.graph_store else {
1493 return Ok(Vec::new());
1494 };
1495
1496 tracing::debug!(
1497 query_len = query.len(),
1498 limit,
1499 max_hops,
1500 "graph: starting recall"
1501 );
1502
1503 let results = crate::graph::retrieval::graph_recall(
1504 store,
1505 self.qdrant.as_deref(),
1506 &self.provider,
1507 query,
1508 limit,
1509 max_hops,
1510 at_timestamp,
1511 temporal_decay_rate,
1512 edge_types,
1513 self.hebbian_reinforcement.is_enabled(),
1514 self.hebbian_lr,
1515 self.embed_timeout,
1516 )
1517 .await?;
1518
1519 tracing::debug!(result_count = results.len(), "graph: recall complete");
1520 #[cfg(feature = "profiling")]
1521 tracing::Span::current().record("result_count", results.len());
1522
1523 Ok(results)
1524 }
1525
1526 #[cfg_attr(
1535 feature = "profiling",
1536 tracing::instrument(name = "memory.recall_graph", skip_all, fields(result_count = tracing::field::Empty))
1537 )]
1538 pub async fn recall_graph_activated(
1539 &self,
1540 query: &str,
1541 limit: usize,
1542 params: crate::graph::SpreadingActivationParams,
1543 edge_types: &[crate::graph::EdgeType],
1544 ) -> Result<Vec<crate::graph::activation::ActivatedFact>, MemoryError> {
1545 let Some(store) = &self.graph_store else {
1546 return Ok(Vec::new());
1547 };
1548
1549 tracing::debug!(
1550 query_len = query.len(),
1551 limit,
1552 "spreading activation: starting graph recall"
1553 );
1554
1555 let embeddings = self.qdrant.as_deref();
1556 let results = crate::graph::retrieval::graph_recall_activated(
1557 store,
1558 embeddings,
1559 &self.provider,
1560 query,
1561 limit,
1562 params,
1563 edge_types,
1564 self.hebbian_reinforcement.is_enabled(),
1565 self.hebbian_lr,
1566 self.embed_timeout,
1567 )
1568 .await?;
1569
1570 tracing::debug!(
1571 result_count = results.len(),
1572 "spreading activation: graph recall complete"
1573 );
1574
1575 Ok(results)
1576 }
1577
1578 #[allow(clippy::too_many_arguments, clippy::too_many_lines)] #[cfg_attr(
1611 feature = "profiling",
1612 tracing::instrument(
1613 name = "memory.recall.graph_view",
1614 skip_all,
1615 fields(view = ?view, result_count = tracing::field::Empty)
1616 )
1617 )]
1618 pub async fn recall_graph_view(
1619 &self,
1620 query: &str,
1621 limit: usize,
1622 view: crate::recall_view::RecallView,
1623 neighbor_cap: usize,
1624 bfs_max_hops: u32,
1625 temporal_decay_rate: f64,
1626 edge_types: &[crate::graph::EdgeType],
1627 sa_params: Option<crate::graph::SpreadingActivationParams>,
1628 ) -> Result<Vec<crate::recall_view::RecalledFact>, MemoryError> {
1629 use crate::recall_view::{RecallView, RecalledFact};
1630
1631 let mut recalled: Vec<RecalledFact> = if let Some(params) = sa_params {
1633 let activated = self
1634 .recall_graph_activated(query, limit, params, edge_types)
1635 .await?;
1636 activated
1637 .into_iter()
1638 .map(|af| {
1639 let activation_score = af.activation_score;
1642 let edge = &af.edge;
1643 let fact = crate::graph::types::GraphFact {
1644 entity_name: String::new(), relation: edge.canonical_relation.clone(),
1646 target_name: String::new(),
1647 fact: edge.fact.clone(),
1648 entity_match_score: activation_score,
1649 hop_distance: 0,
1650 confidence: edge.confidence,
1651 valid_from: if edge.valid_from.is_empty() {
1652 None
1653 } else {
1654 Some(edge.valid_from.clone())
1655 },
1656 edge_type: edge.edge_type,
1657 retrieval_count: edge.retrieval_count,
1658 edge_id: Some(edge.id),
1659 };
1660 RecalledFact {
1661 fact,
1662 activation_score: Some(activation_score),
1663 provenance_message_id: edge.source_message_id,
1664 provenance_snippet: None,
1665 neighbors: Vec::new(),
1666 }
1667 })
1668 .collect()
1669 } else {
1670 let facts = self
1671 .recall_graph(
1672 query,
1673 limit,
1674 bfs_max_hops,
1675 None,
1676 temporal_decay_rate,
1677 edge_types,
1678 )
1679 .await?;
1680 facts
1681 .into_iter()
1682 .map(RecalledFact::from_graph_fact)
1683 .collect()
1684 };
1685
1686 if view == RecallView::Head {
1688 #[cfg(feature = "profiling")]
1689 tracing::Span::current().record("result_count", recalled.len());
1690 return Ok(recalled);
1691 }
1692
1693 if matches!(view, RecallView::ZoomIn | RecallView::ZoomOut) {
1695 let edge_ids: Vec<i64> = recalled.iter().filter_map(|r| r.fact.edge_id).collect();
1696
1697 if !edge_ids.is_empty()
1698 && let Some(ref store) = self.graph_store
1699 {
1700 const MAX_IDS: usize = 490;
1702 let mut edge_to_msg: std::collections::HashMap<i64, MessageId> =
1703 std::collections::HashMap::new();
1704 for chunk in edge_ids.chunks(MAX_IDS) {
1705 match store.source_message_ids_for_edges(chunk).await {
1706 Ok(pairs) => {
1707 for (eid, mid) in pairs {
1708 edge_to_msg.insert(eid, mid);
1709 }
1710 }
1711 Err(e) => {
1712 tracing::warn!(error = %e, "recall_graph_view: provenance fetch failed");
1713 }
1714 }
1715 }
1716
1717 for rf in &mut recalled {
1719 if rf.provenance_message_id.is_none()
1720 && let Some(eid) = rf.fact.edge_id
1721 {
1722 rf.provenance_message_id = edge_to_msg.get(&eid).copied();
1723 }
1724 }
1725
1726 let msg_ids: Vec<MessageId> = recalled
1728 .iter()
1729 .filter_map(|r| r.provenance_message_id)
1730 .collect::<std::collections::HashSet<_>>()
1731 .into_iter()
1732 .collect();
1733
1734 if !msg_ids.is_empty() {
1735 match self.sqlite.messages_by_ids(&msg_ids).await {
1736 Ok(messages) => {
1737 let mut mid_to_snippet: std::collections::HashMap<MessageId, String> =
1738 messages
1739 .into_iter()
1740 .map(|(id, msg)| {
1741 let raw = &msg.content;
1742 let scrubbed: String = raw
1743 .chars()
1744 .map(|c| match c {
1745 '\n' | '\r' | '<' | '>' => ' ',
1746 other => other,
1747 })
1748 .take(200)
1749 .collect();
1750 (id, scrubbed)
1751 })
1752 .collect();
1753 for rf in &mut recalled {
1754 if let Some(mid) = rf.provenance_message_id {
1755 rf.provenance_snippet = mid_to_snippet.remove(&mid);
1756 }
1757 }
1758 }
1759 Err(e) => {
1760 tracing::warn!(error = %e, "recall_graph_view: message snippet fetch failed");
1761 }
1762 }
1763 }
1764 }
1765 }
1766
1767 if view == RecallView::ZoomOut
1769 && let Some(ref store) = self.graph_store
1770 {
1771 type DedupeKey = (String, String, String, crate::graph::EdgeType);
1775 let make_key = |f: &crate::graph::types::GraphFact| -> DedupeKey {
1776 if f.entity_name.is_empty() || f.target_name.is_empty() {
1777 (
1778 f.fact.clone(),
1779 f.relation.clone(),
1780 String::new(),
1781 f.edge_type,
1782 )
1783 } else {
1784 (
1785 f.entity_name.clone(),
1786 f.relation.clone(),
1787 f.target_name.clone(),
1788 f.edge_type,
1789 )
1790 }
1791 };
1792 let mut seen: std::collections::HashSet<DedupeKey> =
1793 recalled.iter().map(|r| make_key(&r.fact)).collect();
1794
1795 let total_neighbor_cap = limit * neighbor_cap;
1796 let mut total_neighbors = 0usize;
1797
1798 for rf in &mut recalled {
1799 if total_neighbors >= total_neighbor_cap {
1800 break;
1801 }
1802 let source_entity_id = match rf.fact.edge_id {
1805 Some(eid) => match store.source_entity_id_for_edge(eid).await {
1806 Ok(Some(id)) => id,
1807 _ => continue,
1808 },
1809 None => continue,
1810 };
1811
1812 let neighbors = match store
1813 .bfs_edges_at_depth(source_entity_id, 1, edge_types)
1814 .await
1815 {
1816 Ok(edges) => edges,
1817 Err(e) => {
1818 tracing::warn!(error = %e, "recall_graph_view: zoom_out bfs failed");
1819 continue;
1820 }
1821 };
1822
1823 let mut added = 0usize;
1824 for n_edge in neighbors {
1825 if added >= neighbor_cap || total_neighbors >= total_neighbor_cap {
1826 break;
1827 }
1828 let key = make_key(&n_edge.fact);
1829 if seen.insert(key) {
1830 rf.neighbors.push(n_edge.fact);
1831 added += 1;
1832 total_neighbors += 1;
1833 }
1834 }
1835 }
1836 }
1837
1838 #[cfg(feature = "profiling")]
1839 tracing::Span::current().record("result_count", recalled.len());
1840 Ok(recalled)
1841 }
1842
1843 pub async fn recall_graph_astar(
1851 &self,
1852 query: &str,
1853 limit: usize,
1854 max_hops: u32,
1855 temporal_decay_rate: f64,
1856 edge_types: &[crate::graph::EdgeType],
1857 ) -> Result<Vec<crate::graph::types::GraphFact>, MemoryError> {
1858 let Some(store) = &self.graph_store else {
1859 return Ok(Vec::new());
1860 };
1861 crate::graph::retrieval_astar::graph_recall_astar(
1862 store,
1863 self.qdrant.as_deref(),
1864 &self.provider,
1865 query,
1866 limit,
1867 max_hops,
1868 edge_types,
1869 temporal_decay_rate,
1870 self.hebbian_reinforcement.is_enabled(),
1871 self.hebbian_lr,
1872 self.query_sensitive_cost,
1873 self.embed_timeout,
1874 )
1875 .await
1876 }
1877
1878 pub async fn recall_graph_watercircles(
1886 &self,
1887 query: &str,
1888 limit: usize,
1889 max_hops: u32,
1890 ring_limit: usize,
1891 temporal_decay_rate: f64,
1892 edge_types: &[crate::graph::EdgeType],
1893 ) -> Result<Vec<crate::graph::types::GraphFact>, MemoryError> {
1894 let Some(store) = &self.graph_store else {
1895 return Ok(Vec::new());
1896 };
1897 crate::graph::retrieval_watercircles::graph_recall_watercircles(
1898 store,
1899 self.qdrant.as_deref(),
1900 &self.provider,
1901 query,
1902 limit,
1903 max_hops,
1904 ring_limit,
1905 edge_types,
1906 temporal_decay_rate,
1907 self.hebbian_reinforcement.is_enabled(),
1908 self.hebbian_lr,
1909 self.embed_timeout,
1910 )
1911 .await
1912 }
1913
1914 pub async fn recall_graph_beam(
1922 &self,
1923 query: &str,
1924 limit: usize,
1925 beam_width: usize,
1926 max_hops: u32,
1927 temporal_decay_rate: f64,
1928 edge_types: &[crate::graph::EdgeType],
1929 ) -> Result<Vec<crate::graph::types::GraphFact>, MemoryError> {
1930 let Some(store) = &self.graph_store else {
1931 return Ok(Vec::new());
1932 };
1933 crate::graph::retrieval_beam::graph_recall_beam(
1934 store,
1935 self.qdrant.as_deref(),
1936 &self.provider,
1937 query,
1938 limit,
1939 beam_width,
1940 max_hops,
1941 edge_types,
1942 temporal_decay_rate,
1943 self.hebbian_reinforcement.is_enabled(),
1944 self.hebbian_lr,
1945 self.embed_timeout,
1946 )
1947 .await
1948 }
1949
1950 pub async fn classify_graph_strategy(&self, query: &str) -> String {
1955 crate::graph::strategy_classifier::classify_retrieval_strategy(&self.provider, query).await
1956 }
1957
1958 #[cfg_attr(
1972 feature = "profiling",
1973 tracing::instrument(
1974 name = "memory.recall_graph_hela",
1975 skip_all,
1976 fields(result_count = tracing::field::Empty)
1977 )
1978 )]
1979 pub async fn recall_graph_hela(
1980 &self,
1981 query: &str,
1982 limit: usize,
1983 params: crate::graph::HelaSpreadParams,
1984 ) -> Result<Vec<crate::graph::HelaFact>, MemoryError> {
1985 let Some(store) = &self.graph_store else {
1986 return Ok(Vec::new());
1987 };
1988 let Some(embeddings) = &self.qdrant else {
1989 return Ok(Vec::new());
1990 };
1991
1992 let store = Arc::clone(store);
1993 let embeddings = Arc::clone(embeddings);
1994 let provider = self.provider.clone();
1995 let hebbian_enabled = self.hebbian_reinforcement.is_enabled();
1996 let hebbian_lr = self.hebbian_lr;
1997
1998 let outer_timeout = hela_outer_timeout(¶ms);
2003
2004 let results = tokio::time::timeout(
2005 outer_timeout,
2006 crate::graph::hela_spreading_recall(
2007 &store,
2008 &embeddings,
2009 &provider,
2010 query,
2011 limit,
2012 ¶ms,
2013 hebbian_enabled,
2014 hebbian_lr,
2015 ),
2016 )
2017 .await
2018 .unwrap_or_else(|_| {
2019 tracing::warn!(
2020 outer_timeout_ms = outer_timeout.as_millis(),
2021 "memory.recall_graph_hela: outer timeout exceeded"
2022 );
2023 Ok(Vec::new())
2024 })?;
2025
2026 #[cfg(feature = "profiling")]
2027 tracing::Span::current().record("result_count", results.len());
2028
2029 Ok(results)
2030 }
2031
2032 async fn batch_increment_access_count(
2040 &self,
2041 message_ids: Vec<MessageId>,
2042 ) -> Result<(), MemoryError> {
2043 if message_ids.is_empty() {
2044 return Ok(());
2045 }
2046 self.sqlite.increment_access_counts(&message_ids).await
2047 }
2048
2049 pub async fn has_embedding(&self, message_id: MessageId) -> Result<bool, MemoryError> {
2055 match &self.qdrant {
2056 Some(qdrant) => qdrant.has_embedding(message_id).await,
2057 None => Ok(false),
2058 }
2059 }
2060
2061 pub async fn embed_missing(
2077 &self,
2078 progress_tx: Option<tokio::sync::watch::Sender<Option<super::BackfillProgress>>>,
2079 ) -> Result<usize, MemoryError> {
2080 if self.qdrant.is_none() || !self.effective_embed_provider().supports_embeddings() {
2081 return Ok(0);
2082 }
2083
2084 let total = self.sqlite.count_unembedded_messages().await?;
2085 if total == 0 {
2086 return Ok(0);
2087 }
2088
2089 if let Some(tx) = &progress_tx {
2090 let _ = tx.send(Some(super::BackfillProgress { done: 0, total }));
2091 }
2092
2093 let mut done = 0usize;
2094 let mut succeeded = 0usize;
2095
2096 loop {
2097 const BATCH_SIZE: usize = 32;
2098 const BATCH_SIZE_I64: i64 = 32;
2099 let rows: Vec<_> = self
2100 .sqlite
2101 .stream_unembedded_messages(BATCH_SIZE_I64)
2102 .try_collect()
2103 .await?;
2104
2105 if rows.is_empty() {
2106 break;
2107 }
2108
2109 let batch_len = rows.len();
2110
2111 let results: Vec<bool> = futures::stream::iter(rows)
2112 .map(|(msg_id, conv_id, role, content)| async move {
2113 self.embed_and_store_regular(msg_id, conv_id, &role, &content)
2114 })
2115 .buffer_unordered(4)
2116 .collect()
2117 .await;
2118
2119 for ok in &results {
2120 done += 1;
2121 if *ok {
2122 succeeded += 1;
2123 }
2124 if let Some(tx) = &progress_tx {
2125 let _ = tx.send(Some(super::BackfillProgress { done, total }));
2126 }
2127 }
2128
2129 let batch_succeeded = results.iter().filter(|&&b| b).count();
2130 if batch_succeeded > 0 {
2131 tracing::debug!("Backfill batch: {batch_succeeded}/{batch_len} embedded");
2132 }
2133
2134 if batch_len < BATCH_SIZE {
2135 break;
2136 }
2137 }
2138
2139 if let Some(tx) = &progress_tx {
2140 let _ = tx.send(None);
2141 }
2142
2143 if done > 0 {
2144 tracing::info!("Embedded {succeeded}/{total} missing messages");
2145 }
2146 Ok(succeeded)
2147 }
2148}
2149
2150#[cfg(test)]
2151mod tests {
2152 use super::*;
2153
2154 #[test]
2159 fn hela_outer_timeout_scales_with_spread_depth() {
2160 let step_budget = std::time::Duration::from_millis(80);
2161 let embed_timeout = std::time::Duration::from_secs(5);
2162 let margin = std::time::Duration::from_millis(250);
2163
2164 for depth in 1..=6u32 {
2165 let params = crate::graph::HelaSpreadParams {
2166 spread_depth: depth,
2167 step_budget: Some(step_budget),
2168 embed_timeout: Some(embed_timeout),
2169 ..Default::default()
2170 };
2171 let expected = embed_timeout + step_budget * (depth + 2) + margin;
2172 assert_eq!(
2173 hela_outer_timeout(¶ms),
2174 expected,
2175 "outer timeout must scale with spread_depth={depth}"
2176 );
2177 }
2178 }
2179
2180 #[test]
2184 fn hela_outer_timeout_clamps_spread_depth_above_six() {
2185 let step_budget = std::time::Duration::from_millis(80);
2186 let embed_timeout = std::time::Duration::from_secs(5);
2187 let params_over = crate::graph::HelaSpreadParams {
2188 spread_depth: 50,
2189 step_budget: Some(step_budget),
2190 embed_timeout: Some(embed_timeout),
2191 ..Default::default()
2192 };
2193 let params_clamped = crate::graph::HelaSpreadParams {
2194 spread_depth: 6,
2195 step_budget: Some(step_budget),
2196 embed_timeout: Some(embed_timeout),
2197 ..Default::default()
2198 };
2199 assert_eq!(
2200 hela_outer_timeout(¶ms_over),
2201 hela_outer_timeout(¶ms_clamped),
2202 "spread_depth above 6 must clamp identically to the algorithm's own [1, 6] bound"
2203 );
2204 }
2205
2206 #[test]
2210 fn hela_outer_timeout_falls_back_when_params_disabled() {
2211 let params = crate::graph::HelaSpreadParams {
2212 spread_depth: 2,
2213 step_budget: None,
2214 embed_timeout: None,
2215 ..Default::default()
2216 };
2217 let expected = std::time::Duration::from_secs(5)
2218 + std::time::Duration::from_millis(80) * 4
2219 + std::time::Duration::from_millis(250);
2220 assert_eq!(hela_outer_timeout(¶ms), expected);
2221 }
2222
2223 #[test]
2224 fn embed_context_default_all_none() {
2225 let ctx = EmbedContext::default();
2226 assert!(ctx.tool_name.is_none());
2227 assert!(ctx.exit_code.is_none());
2228 assert!(ctx.timestamp.is_none());
2229 }
2230
2231 #[test]
2232 fn embed_context_fields_set_correctly() {
2233 let ctx = EmbedContext {
2234 tool_name: Some("shell".to_string()),
2235 exit_code: Some(0),
2236 timestamp: Some("2026-04-04T00:00:00Z".to_string()),
2237 };
2238 assert_eq!(ctx.tool_name.as_deref(), Some("shell"));
2239 assert_eq!(ctx.exit_code, Some(0));
2240 assert_eq!(ctx.timestamp.as_deref(), Some("2026-04-04T00:00:00Z"));
2241 }
2242
2243 #[test]
2244 fn embed_context_non_zero_exit_code() {
2245 let ctx = EmbedContext {
2246 tool_name: Some("shell".to_string()),
2247 exit_code: Some(1),
2248 timestamp: None,
2249 };
2250 assert_eq!(ctx.exit_code, Some(1));
2251 assert!(ctx.timestamp.is_none());
2252 }
2253
2254 async fn make_semantic_memory() -> crate::semantic::SemanticMemory {
2255 let sqlite = crate::store::SqliteStore::new(":memory:").await.unwrap();
2256 make_semantic_memory_with_sqlite(sqlite)
2257 }
2258
2259 fn make_semantic_memory_with_sqlite(
2267 sqlite: crate::store::SqliteStore,
2268 ) -> crate::semantic::SemanticMemory {
2269 use std::sync::Arc;
2270 use std::sync::atomic::AtomicU64;
2271 use zeph_llm::any::AnyProvider;
2272 use zeph_llm::mock::MockProvider;
2273
2274 let provider = AnyProvider::Mock(MockProvider::default());
2275 crate::semantic::SemanticMemory {
2276 sqlite,
2277 qdrant: None,
2278 provider,
2279 embed_provider: None,
2280 embedding_model: "test-model".into(),
2281 vector_weight: 0.7,
2282 keyword_weight: 0.3,
2283 temporal_decay: crate::semantic::TemporalDecay::Disabled,
2284 temporal_decay_half_life_days: 30,
2285 mmr_reranking: crate::semantic::MmrReranking::Disabled,
2286 mmr_lambda: 0.7,
2287 importance_scoring: crate::semantic::ImportanceScoring::Disabled,
2288 importance_weight: 0.15,
2289 token_counter: Arc::new(crate::token_counter::TokenCounter::new()),
2290 graph_store: None,
2291 experience: None,
2292 community_detection_failures: Arc::new(AtomicU64::new(0)),
2293 graph_extraction_count: Arc::new(AtomicU64::new(0)),
2294 graph_extraction_failures: Arc::new(AtomicU64::new(0)),
2295 last_qdrant_warn: Arc::new(AtomicU64::new(0)),
2296 tier_boost_semantic: 1.3,
2297 admission_control: None,
2298 quality_gate: None,
2299 key_facts_dedup_threshold: 0.95,
2300 embed_tasks: std::sync::Mutex::new(tokio::task::JoinSet::new()),
2301 retrieval_depth: 0,
2302 search_prompt_template: String::new(),
2303 depth_below_limit_warned: Arc::new(std::sync::atomic::AtomicBool::new(false)),
2304 missing_placeholder_warned: Arc::new(std::sync::atomic::AtomicBool::new(false)),
2305 reasoning: None,
2306 query_bias_correction: crate::semantic::QueryBiasCorrection::Disabled,
2307 query_bias_profile_weight: 0.25,
2308 profile_centroid: tokio::sync::RwLock::new(None),
2309 profile_centroid_ttl_secs: 300,
2310 hebbian_reinforcement: crate::semantic::HebbianReinforcement::Disabled,
2311 hebbian_lr: 0.1,
2312 hebbian_spread: crate::HelaSpreadRuntime::default(),
2313 retrieval_failure_logger: None,
2314 summarization_llm_timeout_secs: 60,
2315 query_sensitive_cost: false,
2316 five_signal: None,
2317 embed_timeout: std::time::Duration::from_secs(5),
2318 graph_cancel: std::sync::Mutex::new(Vec::new()),
2319 }
2320 }
2321
2322 #[tokio::test]
2323 async fn spawn_embed_bg_returns_true_when_capacity_available() {
2324 let memory = make_semantic_memory().await;
2325 let dispatched = memory.spawn_embed_bg(std::future::ready(()));
2326 assert!(
2327 dispatched,
2328 "spawn_embed_bg must return true when a task was successfully spawned"
2329 );
2330 }
2331
2332 #[tokio::test]
2333 async fn spawn_embed_bg_returns_false_at_capacity() {
2334 let memory = make_semantic_memory().await;
2335
2336 {
2338 let mut tasks = memory.embed_tasks.lock().unwrap();
2339 for _ in 0..MAX_EMBED_BG_TASKS {
2340 tasks.spawn(std::future::pending::<()>());
2341 }
2342 }
2343
2344 let dispatched = memory.spawn_embed_bg(std::future::ready(()));
2345 assert!(
2346 !dispatched,
2347 "spawn_embed_bg must return false when the task limit is reached"
2348 );
2349 }
2350
2351 #[test]
2352 fn qdrant_warn_rate_limit_suppresses_within_window() {
2353 assert!(
2355 should_emit_qdrant_warn(0, 100, 10),
2356 "first call must not be suppressed"
2357 );
2358
2359 assert!(
2361 !should_emit_qdrant_warn(100, 105, 10),
2362 "call within 10s window must be suppressed"
2363 );
2364
2365 assert!(
2367 should_emit_qdrant_warn(100, 110, 10),
2368 "call after window expiry must not be suppressed"
2369 );
2370 }
2371
2372 #[cfg(feature = "test-utils")]
2383 #[tokio::test]
2384 #[ignore = "requires Docker"]
2385 async fn apply_five_signal_scoring_decodes_created_at_on_postgres() {
2386 use std::sync::Arc;
2387 use testcontainers::runners::AsyncRunner as _;
2388 use testcontainers_modules::postgres::Postgres;
2389 use zeph_config::memory::FiveSignalConfig;
2390
2391 let image = Postgres::default();
2392 let container = image.start().await.expect("docker must be available");
2393 let host = container.get_host().await.unwrap();
2394 let port = container.get_host_port_ipv4(5432).await.unwrap();
2395 let url = format!("postgres://postgres:postgres@{host}:{port}/postgres");
2396 let pool = zeph_db::DbConfig {
2397 url,
2398 max_connections: 5,
2399 pool_size: 5,
2400 }
2401 .connect()
2402 .await
2403 .expect("failed to connect to PG");
2404
2405 let pg_store = crate::store::SqliteStore::from_pool(pool.clone())
2406 .await
2407 .unwrap();
2408 let cid = pg_store.create_conversation().await.unwrap();
2409
2410 let session_start = 1_700_000_000_i64;
2415 let fresh = pg_store.save_message(cid, "user", "fresh").await.unwrap();
2416 let stale = pg_store.save_message(cid, "user", "stale").await.unwrap();
2417
2418 for (id, epoch) in [(fresh, session_start), (stale, session_start + 20 * 86_400)] {
2419 #[expect(clippy::cast_precision_loss)]
2420 let epoch_f = epoch as f64;
2421 sqlx::query(zeph_db::sql!(
2422 "UPDATE messages SET created_at = to_timestamp(?) WHERE id = ?"
2423 ))
2424 .bind(epoch_f)
2425 .bind(id)
2426 .execute(&pool)
2427 .await
2428 .unwrap();
2429 }
2430
2431 let graph_store = Arc::new(crate::graph::GraphStore::new(pool.clone()));
2432 let config = FiveSignalConfig {
2433 w_recency: 0.0,
2434 w_relevance: 0.0,
2435 w_frequency: 0.0,
2436 w_causal: 0.0,
2437 w_novelty: 1.0,
2438 ..FiveSignalConfig::default()
2439 };
2440 let fs = crate::five_signal::FiveSignalRuntime::new(
2441 config,
2442 pool,
2443 graph_store,
2444 None,
2445 session_start,
2446 "sess-novelty-test",
2447 );
2448
2449 let memory = make_semantic_memory_with_sqlite(pg_store);
2455 let mut ranked = vec![(fresh, 0.0), (stale, 0.0)];
2456 memory
2457 .apply_five_signal_scoring(&mut ranked, &fs, None)
2458 .await;
2459
2460 assert_eq!(
2461 ranked[0].0, fresh,
2462 "fresher message (created_at == session_start) must rank first by novelty"
2463 );
2464 assert!(
2465 (ranked[0].1 - 1.0).abs() < 1e-6,
2466 "message created at session_start must have novelty ~1.0, got {}",
2467 ranked[0].1
2468 );
2469 assert!(
2470 ranked[1].1 < ranked[0].1,
2471 "message created 20 days later must have strictly lower novelty"
2472 );
2473 }
2474
2475 #[test]
2476 fn qdrant_warn_rate_limit_shared_across_concurrent_sites() {
2477 let shared = Arc::new(AtomicU64::new(0));
2482
2483 let site_a = Arc::clone(&shared);
2484 let site_b = Arc::clone(&shared);
2485
2486 let now_a = 100u64;
2487 let last_a = site_a.load(Ordering::Relaxed);
2488 if should_emit_qdrant_warn(last_a, now_a, QDRANT_WARN_WINDOW_SECS) {
2489 site_a.store(now_a, Ordering::Relaxed);
2490 }
2491
2492 let now_b = 105u64;
2493 let last_b = site_b.load(Ordering::Relaxed);
2494 let warn_b = should_emit_qdrant_warn(last_b, now_b, QDRANT_WARN_WINDOW_SECS);
2495 assert!(
2496 !warn_b,
2497 "site B must be suppressed because site A already warned within the window"
2498 );
2499 }
2500
2501 #[test]
2502 fn warn_qdrant_ensure_failure_updates_shared_atomic_once() {
2503 let shared = Arc::new(AtomicU64::new(0));
2504 let err = MemoryError::InvalidInput("boom".into());
2505
2506 warn_qdrant_ensure_failure(&shared, "site A", &err);
2507 let after_first = shared.load(Ordering::Relaxed);
2508 assert!(after_first > 0, "first call must record a warn timestamp");
2509
2510 warn_qdrant_ensure_failure(&shared, "site B", &err);
2513 assert_eq!(
2514 shared.load(Ordering::Relaxed),
2515 after_first,
2516 "suppressed call must not overwrite the shared warn timestamp"
2517 );
2518 }
2519}