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 let recent_embeddings = self
478 .fetch_recent_embeddings(conversation_id, gate.config().recent_window)
479 .await;
480 if gate
481 .evaluate(content, self.effective_embed_provider(), &recent_embeddings)
482 .await
483 .is_none()
484 {
485 return false;
486 }
487 if let Some(decision) = admission_decision {
488 self.record_admission_sample(conversation_id, role, content, decision, None)
489 .await;
490 }
491 true
492 }
493
494 async fn fetch_recent_embeddings(
504 &self,
505 conversation_id: ConversationId,
506 limit: usize,
507 ) -> Vec<Vec<f32>> {
508 let Some(qdrant) = &self.qdrant else {
509 return Vec::new();
510 };
511 if limit == 0 {
512 return Vec::new();
513 }
514 let limit_u32 = u32::try_from(limit).unwrap_or(u32::MAX);
515 let recent_messages = match self.sqlite.load_history(conversation_id, limit_u32).await {
516 Ok(messages) => messages,
517 Err(e) => {
518 tracing::warn!("quality_gate: failed to load recent history: {e:#}");
519 return Vec::new();
520 }
521 };
522 let ids: Vec<MessageId> = recent_messages
523 .iter()
524 .filter_map(|m| m.metadata.db_id)
525 .map(MessageId)
526 .collect();
527 if ids.is_empty() {
528 return Vec::new();
529 }
530 match qdrant.get_vectors(&ids).await {
531 Ok(vec_map) => vec_map.into_values().collect(),
532 Err(e) => {
533 tracing::warn!("quality_gate: failed to fetch recent embeddings: {e:#}");
534 Vec::new()
535 }
536 }
537 }
538
539 async fn record_admission_sample_opt(
542 &self,
543 conversation_id: ConversationId,
544 role: &str,
545 content: &str,
546 decision: Option<&AdmissionDecision>,
547 message_id: Option<MessageId>,
548 ) {
549 if let Some(decision) = decision {
550 self.record_admission_sample(conversation_id, role, content, decision, message_id)
551 .await;
552 }
553 }
554
555 async fn record_admission_sample(
563 &self,
564 conversation_id: ConversationId,
565 role: &str,
566 content: &str,
567 decision: &AdmissionDecision,
568 message_id: Option<MessageId>,
569 ) {
570 let features_json = match serde_json::to_string(&decision.factors) {
571 Ok(json) => json,
572 Err(e) => {
573 tracing::debug!(error = %e, "admission training: failed to serialize factors");
574 return;
575 }
576 };
577 if let Err(e) = self
578 .sqlite
579 .record_admission_training(AdmissionTrainingInput {
580 message_id,
581 conversation_id,
582 content,
583 role,
584 composite_score: decision.composite_score,
585 was_admitted: decision.admitted,
586 features_json: &features_json,
587 })
588 .await
589 {
590 tracing::debug!(error = %e, "admission training: failed to record sample (non-fatal)");
591 }
592 }
593
594 pub async fn recall_with_category(
602 &self,
603 query: &str,
604 limit: usize,
605 filter: Option<SearchFilter>,
606 category: Option<&str>,
607 ) -> Result<Vec<RecalledMessage>, MemoryError> {
608 let filter_with_category = filter.map(|mut f| {
609 f.category = category.map(str::to_owned);
610 f
611 });
612 self.recall(query, limit, filter_with_category).await
613 }
614
615 pub fn reap_embed_tasks(&self) {
619 if let Ok(mut tasks) = self.embed_tasks.lock() {
620 while tasks.try_join_next().is_some() {}
621 }
622 }
623
624 fn spawn_embed_bg<F>(&self, fut: F) -> bool
628 where
629 F: std::future::Future<Output = ()> + Send + 'static,
630 {
631 let Ok(mut tasks) = self.embed_tasks.lock() else {
632 return false;
633 };
634 while tasks.try_join_next().is_some() {}
636 if tasks.len() >= MAX_EMBED_BG_TASKS {
637 tracing::debug!("background embed task limit reached, skipping");
638 return false;
639 }
640 tasks.spawn(fut);
641 true
642 }
643
644 fn embed_and_store_with_category(
648 &self,
649 message_id: MessageId,
650 conversation_id: ConversationId,
651 role: &str,
652 content: &str,
653 category: Option<&str>,
654 ) -> bool {
655 let Some(qdrant) = self.qdrant.clone() else {
656 return false;
657 };
658 let embed_provider = self.effective_embed_provider().clone();
659 if !embed_provider.supports_embeddings() {
660 return false;
661 }
662 let store_qdrant = Arc::clone(&qdrant);
663 let embedding_model = self.embedding_model.clone();
664 let role = role.to_owned();
665 let category = category.map(str::to_owned);
666 let store_chunk =
667 move |chunk_index: u32,
668 vector: Vec<f32>|
669 -> Pin<Box<dyn Future<Output = Result<(), MemoryError>> + Send>> {
670 let qdrant = Arc::clone(&store_qdrant);
671 let embedding_model = embedding_model.clone();
672 let role = role.clone();
673 let category = category.clone();
674 Box::pin(async move {
675 qdrant
676 .store_with_category(
677 message_id,
678 conversation_id,
679 &role,
680 vector,
681 MessageKind::Regular,
682 &embedding_model,
683 chunk_index,
684 category.as_deref(),
685 )
686 .await
687 .map(|_| ())
688 })
689 };
690 self.spawn_embed_bg(embed_chunk_and_store_bg(
691 EmbedBgArgs {
692 qdrant,
693 embed_provider,
694 message_id,
695 content: content.to_owned(),
696 last_qdrant_warn: Arc::clone(&self.last_qdrant_warn),
697 },
698 "bg embed_category",
699 store_chunk,
700 ))
701 }
702
703 fn embed_and_store_regular(
707 &self,
708 message_id: MessageId,
709 conversation_id: ConversationId,
710 role: &str,
711 content: &str,
712 ) -> bool {
713 let Some(qdrant) = self.qdrant.clone() else {
714 return false;
715 };
716 let embed_provider = self.effective_embed_provider().clone();
717 if !embed_provider.supports_embeddings() {
718 return false;
719 }
720 let store_qdrant = Arc::clone(&qdrant);
721 let embedding_model = self.embedding_model.clone();
722 let role = role.to_owned();
723 let store_chunk =
724 move |chunk_index: u32,
725 vector: Vec<f32>|
726 -> Pin<Box<dyn Future<Output = Result<(), MemoryError>> + Send>> {
727 let qdrant = Arc::clone(&store_qdrant);
728 let embedding_model = embedding_model.clone();
729 let role = role.clone();
730 Box::pin(async move {
731 qdrant
732 .store(
733 message_id,
734 conversation_id,
735 &role,
736 vector,
737 MessageKind::Regular,
738 &embedding_model,
739 chunk_index,
740 )
741 .await
742 .map(|_| ())
743 })
744 };
745 self.spawn_embed_bg(embed_chunk_and_store_bg(
746 EmbedBgArgs {
747 qdrant,
748 embed_provider,
749 message_id,
750 content: content.to_owned(),
751 last_qdrant_warn: Arc::clone(&self.last_qdrant_warn),
752 },
753 "bg embed_regular",
754 store_chunk,
755 ))
756 }
757
758 fn embed_chunks_with_tool_context(
762 &self,
763 message_id: MessageId,
764 conversation_id: ConversationId,
765 role: &str,
766 content: &str,
767 embed_ctx: EmbedContext,
768 ) -> bool {
769 let Some(qdrant) = self.qdrant.clone() else {
770 return false;
771 };
772 let embed_provider = self.effective_embed_provider().clone();
773 if !embed_provider.supports_embeddings() {
774 return false;
775 }
776 let store_qdrant = Arc::clone(&qdrant);
777 let embedding_model = self.embedding_model.clone();
778 let role = role.to_owned();
779 let store_chunk =
780 move |chunk_index: u32,
781 vector: Vec<f32>|
782 -> Pin<Box<dyn Future<Output = Result<(), MemoryError>> + Send>> {
783 let qdrant = Arc::clone(&store_qdrant);
784 let embedding_model = embedding_model.clone();
785 let role = role.clone();
786 let embed_ctx = embed_ctx.clone();
787 Box::pin(async move {
788 if let Some(tool_name) = embed_ctx.tool_name {
789 qdrant
790 .store_with_tool_context(
791 message_id,
792 conversation_id,
793 &role,
794 vector,
795 MessageKind::Regular,
796 &embedding_model,
797 chunk_index,
798 &tool_name,
799 embed_ctx.exit_code,
800 embed_ctx.timestamp.as_deref(),
801 )
802 .await
803 .map(|_| ())
804 } else {
805 qdrant
806 .store(
807 message_id,
808 conversation_id,
809 &role,
810 vector,
811 MessageKind::Regular,
812 &embedding_model,
813 chunk_index,
814 )
815 .await
816 .map(|_| ())
817 }
818 })
819 };
820 self.spawn_embed_bg(embed_chunk_and_store_bg(
821 EmbedBgArgs {
822 qdrant,
823 embed_provider,
824 message_id,
825 content: content.to_owned(),
826 last_qdrant_warn: Arc::clone(&self.last_qdrant_warn),
827 },
828 "bg embed_tool",
829 store_chunk,
830 ))
831 }
832
833 pub async fn save_only(
841 &self,
842 conversation_id: ConversationId,
843 role: &str,
844 content: &str,
845 parts_json: &str,
846 ) -> Result<MessageId, MemoryError> {
847 self.sqlite
848 .save_message_with_parts(conversation_id, role, content, parts_json)
849 .await
850 }
851
852 #[cfg_attr(
862 feature = "profiling",
863 tracing::instrument(name = "memory.recall", skip_all, fields(query_len = %query.len(), result_count = tracing::field::Empty, top_score = tracing::field::Empty))
864 )]
865 pub async fn recall(
866 &self,
867 query: &str,
868 limit: usize,
869 filter: Option<SearchFilter>,
870 ) -> Result<Vec<RecalledMessage>, MemoryError> {
871 let conversation_id = filter.as_ref().and_then(|f| f.conversation_id);
872
873 tracing::debug!(
874 query_len = query.len(),
875 limit,
876 has_filter = filter.is_some(),
877 conversation_id = conversation_id.map(|c| c.0),
878 has_qdrant = self.qdrant.is_some(),
879 "recall: starting hybrid search"
880 );
881
882 let keyword_results = match self
883 .sqlite
884 .keyword_search(query, self.effective_depth(limit), conversation_id)
885 .await
886 {
887 Ok(results) => results,
888 Err(e) => {
889 tracing::warn!("FTS5 keyword search failed: {e:#}");
890 Vec::new()
891 }
892 };
893
894 let vector_results = if let Some(qdrant) = &self.qdrant
895 && self.effective_embed_provider().supports_embeddings()
896 {
897 let embed_input = self.apply_search_prompt(query);
898 let query_vector = match tokio::time::timeout(
899 self.embed_timeout,
900 self.effective_embed_provider().embed(&embed_input),
901 )
902 .await
903 {
904 Ok(Ok(v)) => v,
905 Ok(Err(e)) => return Err(e.into()),
906 Err(_) => {
907 tracing::warn!("recall_semantic: embed timed out, returning empty results");
908 return Ok(Vec::new());
909 }
910 };
911 let query_vector = self.apply_query_bias(query, query_vector).await;
912 qdrant.ensure_collection_for_vector(&query_vector).await?;
913 qdrant
914 .search(&query_vector, self.effective_depth(limit), filter)
915 .await?
916 } else {
917 Vec::new()
918 };
919
920 let results = self
921 .recall_merge_and_rank(keyword_results, vector_results, limit, None)
922 .await?;
923 #[cfg(feature = "profiling")]
924 {
925 let span = tracing::Span::current();
926 span.record("result_count", results.len());
927 if let Some(top) = results.first() {
928 span.record("top_score", top.score);
929 }
930 }
931 Ok(results)
932 }
933
934 #[cfg_attr(
935 feature = "profiling",
936 tracing::instrument(name = "memory.recall.fts5", skip_all, fields(query_len = %query.len()))
937 )]
938 pub(super) async fn recall_fts5_raw(
939 &self,
940 query: &str,
941 limit: usize,
942 conversation_id: Option<ConversationId>,
943 ) -> Result<Vec<(MessageId, f64)>, MemoryError> {
944 self.sqlite
945 .keyword_search(query, self.effective_depth(limit), conversation_id)
946 .await
947 }
948
949 #[cfg_attr(
950 feature = "profiling",
951 tracing::instrument(name = "memory.recall.vectors", skip_all, fields(query_len = %query.len()))
952 )]
953 pub(super) async fn recall_vectors_raw(
954 &self,
955 query: &str,
956 limit: usize,
957 filter: Option<SearchFilter>,
958 ) -> Result<Vec<crate::embedding_store::SearchResult>, MemoryError> {
959 let Some(qdrant) = &self.qdrant else {
960 return Ok(Vec::new());
961 };
962 if !self.effective_embed_provider().supports_embeddings() {
963 return Ok(Vec::new());
964 }
965 let embed_input = self.apply_search_prompt(query);
966 let query_vector = match tokio::time::timeout(
967 self.embed_timeout,
968 self.effective_embed_provider().embed(&embed_input),
969 )
970 .await
971 {
972 Ok(Ok(v)) => v,
973 Ok(Err(e)) => return Err(e.into()),
974 Err(_) => {
975 tracing::warn!("recall_vectors_raw: embed timed out, returning empty results");
976 return Ok(Vec::new());
977 }
978 };
979 let query_vector = self.apply_query_bias(query, query_vector).await;
980 qdrant.ensure_collection_for_vector(&query_vector).await?;
981 qdrant
982 .search(&query_vector, self.effective_depth(limit), filter)
983 .await
984 }
985
986 #[cfg_attr(
995 feature = "profiling",
996 tracing::instrument(name = "memory.recall.merge_and_rank", skip_all, fields(kw_count = keyword_results.len(), vec_count = vector_results.len()))
997 )]
998 #[allow(clippy::cast_possible_truncation, clippy::too_many_lines)]
999 pub(super) async fn recall_merge_and_rank(
1000 &self,
1001 keyword_results: Vec<(MessageId, f64)>,
1002 vector_results: Vec<crate::embedding_store::SearchResult>,
1003 limit: usize,
1004 goal_entity_id: Option<i64>,
1005 ) -> Result<Vec<RecalledMessage>, MemoryError> {
1006 tracing::debug!(
1007 vector_count = vector_results.len(),
1008 keyword_count = keyword_results.len(),
1009 limit,
1010 "recall: merging search results"
1011 );
1012
1013 let mut scores: std::collections::HashMap<MessageId, f64> =
1014 std::collections::HashMap::new();
1015
1016 if !vector_results.is_empty() {
1017 let max_vs = vector_results
1018 .iter()
1019 .map(|r| r.score)
1020 .fold(f32::NEG_INFINITY, f32::max);
1021 let norm = if max_vs > 0.0 { max_vs } else { 1.0 };
1022 for r in &vector_results {
1023 let normalized = f64::from(r.score / norm);
1024 *scores.entry(r.message_id).or_default() += normalized * self.vector_weight;
1025 }
1026 }
1027
1028 if !keyword_results.is_empty() {
1029 let max_ks = keyword_results
1030 .iter()
1031 .map(|r| r.1)
1032 .fold(f64::NEG_INFINITY, f64::max);
1033 let norm = if max_ks > 0.0 { max_ks } else { 1.0 };
1034 for &(msg_id, score) in &keyword_results {
1035 let normalized = score / norm;
1036 *scores.entry(msg_id).or_default() += normalized * self.keyword_weight;
1037 }
1038 }
1039
1040 if scores.is_empty() {
1041 tracing::debug!("recall: empty merge, no overlapping scores");
1042 return Ok(Vec::new());
1043 }
1044
1045 let mut ranked: Vec<(MessageId, f64)> = scores.into_iter().collect();
1046 ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
1047
1048 tracing::debug!(
1049 merged = ranked.len(),
1050 top_score = ranked.first().map(|r| r.1),
1051 bottom_score = ranked.last().map(|r| r.1),
1052 vector_weight = %self.vector_weight,
1053 keyword_weight = %self.keyword_weight,
1054 "recall: weighted merge complete"
1055 );
1056
1057 if self.temporal_decay.is_enabled() && self.temporal_decay_half_life_days > 0 {
1058 let ids: Vec<MessageId> = ranked.iter().map(|r| r.0).collect();
1059 match self.sqlite.message_timestamps(&ids).await {
1060 Ok(timestamps) => {
1061 apply_temporal_decay(
1062 &mut ranked,
1063 ×tamps,
1064 self.temporal_decay_half_life_days,
1065 );
1066 ranked
1067 .sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
1068 tracing::debug!(
1069 half_life_days = self.temporal_decay_half_life_days,
1070 top_score_after = ranked.first().map(|r| r.1),
1071 "recall: temporal decay applied"
1072 );
1073 }
1074 Err(e) => {
1075 tracing::warn!("temporal decay: failed to fetch timestamps: {e:#}");
1076 }
1077 }
1078 }
1079
1080 if self.mmr_reranking.is_enabled() && !vector_results.is_empty() {
1081 if let Some(qdrant) = &self.qdrant {
1082 let ids: Vec<MessageId> = ranked.iter().map(|r| r.0).collect();
1083 match qdrant.get_vectors(&ids).await {
1084 Ok(vec_map) if !vec_map.is_empty() => {
1085 let ranked_len_before = ranked.len();
1086 ranked = apply_mmr(&ranked, &vec_map, self.mmr_lambda, limit);
1087 tracing::debug!(
1088 before = ranked_len_before,
1089 after = ranked.len(),
1090 lambda = %self.mmr_lambda,
1091 "recall: mmr re-ranked"
1092 );
1093 }
1094 Ok(_) => {
1095 ranked.truncate(limit);
1096 }
1097 Err(e) => {
1098 tracing::warn!("MMR: failed to fetch vectors: {e:#}");
1099 ranked.truncate(limit);
1100 }
1101 }
1102 } else {
1103 ranked.truncate(limit);
1104 }
1105 } else {
1106 ranked.truncate(limit);
1107 }
1108
1109 if self.importance_scoring.is_enabled() && !ranked.is_empty() {
1110 let ids: Vec<MessageId> = ranked.iter().map(|r| r.0).collect();
1111 match self.sqlite.fetch_importance_scores(&ids).await {
1112 Ok(scores) => {
1113 for (msg_id, score) in &mut ranked {
1114 if let Some(&imp) = scores.get(msg_id) {
1115 *score += imp * self.importance_weight;
1116 }
1117 }
1118 ranked
1119 .sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
1120 tracing::debug!(
1121 importance_weight = %self.importance_weight,
1122 "recall: importance scores blended"
1123 );
1124 }
1125 Err(e) => {
1126 tracing::warn!("importance scoring: failed to fetch scores: {e:#}");
1127 }
1128 }
1129 }
1130
1131 if (self.tier_boost_semantic - 1.0).abs() > f64::EPSILON && !ranked.is_empty() {
1135 let ids: Vec<MessageId> = ranked.iter().map(|r| r.0).collect();
1136 match self.sqlite.fetch_tiers(&ids).await {
1137 Ok(tiers) => {
1138 let bonus = self.tier_boost_semantic - 1.0;
1139 let mut boosted = false;
1140 for (msg_id, score) in &mut ranked {
1141 if tiers.get(msg_id).map(String::as_str) == Some("semantic") {
1142 *score += bonus;
1143 boosted = true;
1144 }
1145 }
1146 if boosted {
1147 ranked.sort_by(|a, b| {
1148 b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)
1149 });
1150 tracing::debug!(
1151 tier_boost = %self.tier_boost_semantic,
1152 "recall: semantic tier boost applied"
1153 );
1154 }
1155 }
1156 Err(e) => {
1157 tracing::warn!("tier boost: failed to fetch tiers: {e:#}");
1158 }
1159 }
1160 }
1161
1162 if let Some(fs) = &self.five_signal
1164 && !fs.weights.is_baseline()
1165 {
1166 self.apply_five_signal_scoring(&mut ranked, fs, goal_entity_id)
1167 .await;
1168 }
1169
1170 let ids: Vec<MessageId> = ranked.iter().map(|r| r.0).collect();
1171
1172 if let Some(fs) = &self.five_signal {
1174 for id in &ids {
1175 fs.access_cache
1176 .log_access(*id, "message", &fs.session_id)
1177 .await;
1178 }
1179 fs.metrics.inc_recall();
1180 }
1181
1182 if !ids.is_empty()
1183 && let Err(e) = self.batch_increment_access_count(ids.clone()).await
1184 {
1185 tracing::warn!("recall: failed to increment access counts: {e:#}");
1186 }
1187
1188 if let Err(e) = self.sqlite.mark_training_recalled(&ids).await {
1190 tracing::debug!(
1191 error = %e,
1192 "recall: failed to mark training data as recalled (non-fatal)"
1193 );
1194 }
1195
1196 let messages = self.sqlite.messages_by_ids(&ids).await?;
1197 let msg_map: std::collections::HashMap<MessageId, _> = messages.into_iter().collect();
1198
1199 let recalled: Vec<RecalledMessage> = ranked
1200 .iter()
1201 .filter_map(|(msg_id, score)| {
1202 msg_map.get(msg_id).map(|msg| RecalledMessage {
1203 message: msg.clone(),
1204 #[expect(clippy::cast_possible_truncation)]
1205 score: *score as f32,
1206 })
1207 })
1208 .collect();
1209
1210 tracing::debug!(final_count = recalled.len(), "recall: final results");
1211
1212 Ok(recalled)
1213 }
1214
1215 async fn apply_five_signal_scoring(
1221 &self,
1222 ranked: &mut [(MessageId, f64)],
1223 fs: &crate::five_signal::FiveSignalRuntime,
1224 goal_entity_id: Option<i64>,
1225 ) {
1226 use crate::five_signal::causal_distance::CausalDistanceComputer;
1227 use crate::five_signal::scoring::{CandidateSignals, apply_five_signal_scoring};
1228 use sqlx::Row as _;
1229
1230 let ids: Vec<MessageId> = ranked.iter().map(|r| r.0).collect();
1231
1232 let freq_map = match fs
1234 .access_cache
1235 .load_for_candidates(&fs.session_id, &ids)
1236 .await
1237 {
1238 Ok(m) => m,
1239 Err(e) => {
1240 tracing::warn!(error = %e, "five_signal: failed to load access frequencies (skipping)");
1241 return;
1242 }
1243 };
1244
1245 let created_at_map: std::collections::HashMap<MessageId, i64> = {
1247 let id_vals: Vec<i64> = ids.iter().map(|id| id.0).collect();
1248 let placeholders = zeph_db::placeholder_list(1, id_vals.len());
1249 let created_at_epoch =
1250 <zeph_db::ActiveDialect as zeph_db::dialect::Dialect>::epoch_from_col("created_at");
1251 let sql = format!(
1252 "SELECT id, {created_at_epoch} AS created_at FROM messages \
1253 WHERE id IN ({placeholders}) AND deleted_at IS NULL"
1254 );
1255 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1256 for id in &id_vals {
1257 q = q.bind(id);
1258 }
1259 match q.fetch_all(&fs.pool).await {
1260 Ok(rows) => rows
1261 .iter()
1262 .map(|row| {
1263 (
1264 MessageId(row.get::<i64, _>("id")),
1265 row.get::<i64, _>("created_at"),
1266 )
1267 })
1268 .collect(),
1269 Err(e) => {
1270 tracing::warn!(error = %e, "five_signal: failed to fetch created_at (skipping novelty)");
1271 std::collections::HashMap::new()
1272 }
1273 }
1274 };
1275
1276 let causal_distance_map: std::collections::HashMap<i64, u32> = {
1280 let entity_ids: Vec<i64> = ids.iter().map(|id| id.0).collect();
1281 match fs
1282 .causal_computer
1283 .compute(goal_entity_id, &entity_ids)
1284 .await
1285 {
1286 Ok(m) => m,
1287 Err(e) => {
1288 tracing::warn!(error = %e, "five_signal: causal BFS failed (using neutral)");
1289 std::collections::HashMap::new()
1290 }
1291 }
1292 };
1293 let neutral_causal_score =
1294 CausalDistanceComputer::distance_to_score(fs.config.neutral_causal_distance);
1295
1296 let mut signals_map = std::collections::HashMap::with_capacity(ids.len());
1297 for &(msg_id, base_score) in ranked.iter() {
1298 let frequency = freq_map.get(&msg_id).copied().unwrap_or(0.0);
1299 let half = base_score / 2.0;
1302 let fact_created_at = created_at_map
1303 .get(&msg_id)
1304 .copied()
1305 .unwrap_or(fs.session_start);
1306 let novelty = fs.novelty_computer.compute(fact_created_at);
1307 let causal = causal_distance_map
1308 .get(&msg_id.0)
1309 .map_or(neutral_causal_score, |&d| {
1310 CausalDistanceComputer::distance_to_score(d)
1311 });
1312 signals_map.insert(
1313 msg_id,
1314 CandidateSignals {
1315 recency: half,
1316 relevance: half,
1317 frequency,
1318 causal,
1319 novelty,
1320 },
1321 );
1322 }
1323
1324 apply_five_signal_scoring(ranked, &fs.weights, &signals_map);
1325
1326 tracing::debug!(
1327 candidate_count = ids.len(),
1328 "recall: five-signal scoring applied"
1329 );
1330 }
1331
1332 async fn recall_by_route(
1338 &self,
1339 route: crate::router::MemoryRoute,
1340 query: &str,
1341 limit: usize,
1342 filter: Option<crate::embedding_store::SearchFilter>,
1343 ) -> Result<
1344 (
1345 Vec<(crate::types::MessageId, f64)>,
1346 Vec<crate::embedding_store::SearchResult>,
1347 ),
1348 MemoryError,
1349 > {
1350 use crate::router::MemoryRoute;
1351
1352 let conversation_id = filter.as_ref().and_then(|f| f.conversation_id);
1353
1354 let results: (
1355 Vec<(crate::types::MessageId, f64)>,
1356 Vec<crate::embedding_store::SearchResult>,
1357 ) = match route {
1358 MemoryRoute::Keyword => {
1359 let kw = self.recall_fts5_raw(query, limit, conversation_id).await?;
1360 (kw, Vec::new())
1361 }
1362 MemoryRoute::Hybrid => {
1363 let kw = match self.recall_fts5_raw(query, limit, conversation_id).await {
1364 Ok(r) => r,
1365 Err(e) => {
1366 tracing::warn!("FTS5 keyword search failed: {e:#}");
1367 Vec::new()
1368 }
1369 };
1370 let vr = self.recall_vectors_raw(query, limit, filter).await?;
1371 (kw, vr)
1372 }
1373 MemoryRoute::Episodic => {
1382 let range = crate::router::resolve_temporal_range(query, chrono::Utc::now());
1383 let cleaned = crate::router::strip_temporal_keywords(query);
1384 let search_query = if cleaned.is_empty() { query } else { &cleaned };
1385 let kw = if let Some(ref r) = range {
1386 self.sqlite
1387 .keyword_search_with_time_range(
1388 search_query,
1389 limit,
1390 conversation_id,
1391 r.after.as_deref(),
1392 r.before.as_deref(),
1393 )
1394 .await?
1395 } else {
1396 self.recall_fts5_raw(search_query, limit, conversation_id)
1397 .await?
1398 };
1399 tracing::debug!(
1400 has_range = range.is_some(),
1401 cleaned_query = %search_query,
1402 keyword_count = kw.len(),
1403 "recall: episodic path"
1404 );
1405 (kw, Vec::new())
1406 }
1407 MemoryRoute::Graph => {
1410 let kw = match self.recall_fts5_raw(query, limit, conversation_id).await {
1411 Ok(r) => r,
1412 Err(e) => {
1413 tracing::warn!("FTS5 keyword search failed (graph→hybrid fallback): {e:#}");
1414 Vec::new()
1415 }
1416 };
1417 let vr = self.recall_vectors_raw(query, limit, filter).await?;
1418 (kw, vr)
1419 }
1420 _ => {
1421 let vr = self.recall_vectors_raw(query, limit, filter).await?;
1422 (Vec::new(), vr)
1423 }
1424 };
1425 Ok(results)
1426 }
1427
1428 #[cfg_attr(
1440 feature = "profiling",
1441 tracing::instrument(name = "memory.recall", skip_all, fields(query_len = %query.len(), result_count = tracing::field::Empty))
1442 )]
1443 pub async fn recall_routed(
1444 &self,
1445 query: &str,
1446 limit: usize,
1447 filter: Option<SearchFilter>,
1448 router: &dyn crate::router::MemoryRouter,
1449 goal_entity_id: Option<i64>,
1450 ) -> Result<Vec<RecalledMessage>, MemoryError> {
1451 let route = router.route(query);
1452 tracing::debug!(?route, query_len = query.len(), "memory routing decision");
1453
1454 let (keyword_results, vector_results) =
1455 self.recall_by_route(route, query, limit, filter).await?;
1456
1457 tracing::debug!(
1458 keyword_count = keyword_results.len(),
1459 vector_count = vector_results.len(),
1460 "recall: routed search results"
1461 );
1462
1463 self.recall_merge_and_rank(keyword_results, vector_results, limit, goal_entity_id)
1464 .await
1465 }
1466
1467 #[cfg_attr(
1481 feature = "profiling",
1482 tracing::instrument(name = "memory.recall", skip_all, fields(query_len = %query.len(), result_count = tracing::field::Empty))
1483 )]
1484 pub async fn recall_routed_async(
1485 &self,
1486 query: &str,
1487 limit: usize,
1488 filter: Option<crate::embedding_store::SearchFilter>,
1489 router: &dyn crate::router::AsyncMemoryRouter,
1490 goal_entity_id: Option<i64>,
1491 ) -> Result<Vec<RecalledMessage>, MemoryError> {
1492 let decision = router.route_async(query).await;
1493 let route = decision.route;
1494 tracing::debug!(
1495 ?route,
1496 confidence = decision.confidence,
1497 query_len = query.len(),
1498 "memory routing decision (async)"
1499 );
1500
1501 let (keyword_results, vector_results) =
1502 self.recall_by_route(route, query, limit, filter).await?;
1503
1504 tracing::debug!(
1505 keyword_count = keyword_results.len(),
1506 vector_count = vector_results.len(),
1507 "recall: routed search results (async)"
1508 );
1509
1510 self.recall_merge_and_rank(keyword_results, vector_results, limit, goal_entity_id)
1511 .await
1512 }
1513
1514 #[cfg_attr(
1528 feature = "profiling",
1529 tracing::instrument(name = "memory.recall_graph", skip_all, fields(result_count = tracing::field::Empty))
1530 )]
1531 pub async fn recall_graph(
1532 &self,
1533 query: &str,
1534 limit: usize,
1535 max_hops: u32,
1536 at_timestamp: Option<&str>,
1537 temporal_decay_rate: f64,
1538 edge_types: &[crate::graph::EdgeType],
1539 ) -> Result<Vec<crate::graph::types::GraphFact>, MemoryError> {
1540 let Some(store) = &self.graph_store else {
1541 return Ok(Vec::new());
1542 };
1543
1544 tracing::debug!(
1545 query_len = query.len(),
1546 limit,
1547 max_hops,
1548 "graph: starting recall"
1549 );
1550
1551 let results = crate::graph::retrieval::graph_recall(
1552 store,
1553 self.qdrant.as_deref(),
1554 &self.provider,
1555 query,
1556 limit,
1557 max_hops,
1558 at_timestamp,
1559 temporal_decay_rate,
1560 edge_types,
1561 self.hebbian_reinforcement.is_enabled(),
1562 self.hebbian_lr,
1563 self.embed_timeout,
1564 )
1565 .await?;
1566
1567 tracing::debug!(result_count = results.len(), "graph: recall complete");
1568 #[cfg(feature = "profiling")]
1569 tracing::Span::current().record("result_count", results.len());
1570
1571 Ok(results)
1572 }
1573
1574 #[cfg_attr(
1583 feature = "profiling",
1584 tracing::instrument(name = "memory.recall_graph", skip_all, fields(result_count = tracing::field::Empty))
1585 )]
1586 pub async fn recall_graph_activated(
1587 &self,
1588 query: &str,
1589 limit: usize,
1590 params: crate::graph::SpreadingActivationParams,
1591 edge_types: &[crate::graph::EdgeType],
1592 ) -> Result<Vec<crate::graph::activation::ActivatedFact>, MemoryError> {
1593 let Some(store) = &self.graph_store else {
1594 return Ok(Vec::new());
1595 };
1596
1597 tracing::debug!(
1598 query_len = query.len(),
1599 limit,
1600 "spreading activation: starting graph recall"
1601 );
1602
1603 let embeddings = self.qdrant.as_deref();
1604 let results = crate::graph::retrieval::graph_recall_activated(
1605 store,
1606 embeddings,
1607 &self.provider,
1608 query,
1609 limit,
1610 params,
1611 edge_types,
1612 self.hebbian_reinforcement.is_enabled(),
1613 self.hebbian_lr,
1614 self.embed_timeout,
1615 )
1616 .await?;
1617
1618 tracing::debug!(
1619 result_count = results.len(),
1620 "spreading activation: graph recall complete"
1621 );
1622
1623 Ok(results)
1624 }
1625
1626 #[allow(clippy::too_many_arguments, clippy::too_many_lines)] #[cfg_attr(
1659 feature = "profiling",
1660 tracing::instrument(
1661 name = "memory.recall.graph_view",
1662 skip_all,
1663 fields(view = ?view, result_count = tracing::field::Empty)
1664 )
1665 )]
1666 pub async fn recall_graph_view(
1667 &self,
1668 query: &str,
1669 limit: usize,
1670 view: crate::recall_view::RecallView,
1671 neighbor_cap: usize,
1672 bfs_max_hops: u32,
1673 temporal_decay_rate: f64,
1674 edge_types: &[crate::graph::EdgeType],
1675 sa_params: Option<crate::graph::SpreadingActivationParams>,
1676 ) -> Result<Vec<crate::recall_view::RecalledFact>, MemoryError> {
1677 use crate::recall_view::{RecallView, RecalledFact};
1678
1679 let mut recalled: Vec<RecalledFact> = if let Some(params) = sa_params {
1681 let activated = self
1682 .recall_graph_activated(query, limit, params, edge_types)
1683 .await?;
1684 activated
1685 .into_iter()
1686 .map(|af| {
1687 let activation_score = af.activation_score;
1690 let edge = &af.edge;
1691 let fact = crate::graph::types::GraphFact {
1692 entity_name: String::new(), relation: edge.canonical_relation.clone(),
1694 target_name: String::new(),
1695 fact: edge.fact.clone(),
1696 entity_match_score: activation_score,
1697 hop_distance: 0,
1698 confidence: edge.confidence,
1699 valid_from: if edge.valid_from.is_empty() {
1700 None
1701 } else {
1702 Some(edge.valid_from.clone())
1703 },
1704 edge_type: edge.edge_type,
1705 retrieval_count: edge.retrieval_count,
1706 edge_id: Some(edge.id),
1707 };
1708 RecalledFact {
1709 fact,
1710 activation_score: Some(activation_score),
1711 provenance_message_id: edge.source_message_id,
1712 provenance_snippet: None,
1713 neighbors: Vec::new(),
1714 }
1715 })
1716 .collect()
1717 } else {
1718 let facts = self
1719 .recall_graph(
1720 query,
1721 limit,
1722 bfs_max_hops,
1723 None,
1724 temporal_decay_rate,
1725 edge_types,
1726 )
1727 .await?;
1728 facts
1729 .into_iter()
1730 .map(RecalledFact::from_graph_fact)
1731 .collect()
1732 };
1733
1734 if view == RecallView::Head {
1736 #[cfg(feature = "profiling")]
1737 tracing::Span::current().record("result_count", recalled.len());
1738 return Ok(recalled);
1739 }
1740
1741 if matches!(view, RecallView::ZoomIn | RecallView::ZoomOut) {
1743 let edge_ids: Vec<i64> = recalled.iter().filter_map(|r| r.fact.edge_id).collect();
1744
1745 if !edge_ids.is_empty()
1746 && let Some(ref store) = self.graph_store
1747 {
1748 const MAX_IDS: usize = 490;
1750 let mut edge_to_msg: std::collections::HashMap<i64, MessageId> =
1751 std::collections::HashMap::new();
1752 for chunk in edge_ids.chunks(MAX_IDS) {
1753 match store.source_message_ids_for_edges(chunk).await {
1754 Ok(pairs) => {
1755 for (eid, mid) in pairs {
1756 edge_to_msg.insert(eid, mid);
1757 }
1758 }
1759 Err(e) => {
1760 tracing::warn!(error = %e, "recall_graph_view: provenance fetch failed");
1761 }
1762 }
1763 }
1764
1765 for rf in &mut recalled {
1767 if rf.provenance_message_id.is_none()
1768 && let Some(eid) = rf.fact.edge_id
1769 {
1770 rf.provenance_message_id = edge_to_msg.get(&eid).copied();
1771 }
1772 }
1773
1774 let msg_ids: Vec<MessageId> = recalled
1776 .iter()
1777 .filter_map(|r| r.provenance_message_id)
1778 .collect::<std::collections::HashSet<_>>()
1779 .into_iter()
1780 .collect();
1781
1782 if !msg_ids.is_empty() {
1783 match self.sqlite.messages_by_ids(&msg_ids).await {
1784 Ok(messages) => {
1785 let mut mid_to_snippet: std::collections::HashMap<MessageId, String> =
1786 messages
1787 .into_iter()
1788 .map(|(id, msg)| {
1789 let raw = &msg.content;
1790 let scrubbed: String = raw
1791 .chars()
1792 .map(|c| match c {
1793 '\n' | '\r' | '<' | '>' => ' ',
1794 other => other,
1795 })
1796 .take(200)
1797 .collect();
1798 (id, scrubbed)
1799 })
1800 .collect();
1801 for rf in &mut recalled {
1802 if let Some(mid) = rf.provenance_message_id {
1803 rf.provenance_snippet = mid_to_snippet.remove(&mid);
1804 }
1805 }
1806 }
1807 Err(e) => {
1808 tracing::warn!(error = %e, "recall_graph_view: message snippet fetch failed");
1809 }
1810 }
1811 }
1812 }
1813 }
1814
1815 if view == RecallView::ZoomOut
1817 && let Some(ref store) = self.graph_store
1818 {
1819 type DedupeKey = (String, String, String, crate::graph::EdgeType);
1823 let make_key = |f: &crate::graph::types::GraphFact| -> DedupeKey {
1824 if f.entity_name.is_empty() || f.target_name.is_empty() {
1825 (
1826 f.fact.clone(),
1827 f.relation.clone(),
1828 String::new(),
1829 f.edge_type,
1830 )
1831 } else {
1832 (
1833 f.entity_name.clone(),
1834 f.relation.clone(),
1835 f.target_name.clone(),
1836 f.edge_type,
1837 )
1838 }
1839 };
1840 let mut seen: std::collections::HashSet<DedupeKey> =
1841 recalled.iter().map(|r| make_key(&r.fact)).collect();
1842
1843 let total_neighbor_cap = limit * neighbor_cap;
1844 let mut total_neighbors = 0usize;
1845
1846 for rf in &mut recalled {
1847 if total_neighbors >= total_neighbor_cap {
1848 break;
1849 }
1850 let source_entity_id = match rf.fact.edge_id {
1853 Some(eid) => match store.source_entity_id_for_edge(eid).await {
1854 Ok(Some(id)) => id,
1855 _ => continue,
1856 },
1857 None => continue,
1858 };
1859
1860 let neighbors = match store
1861 .bfs_edges_at_depth(source_entity_id, 1, edge_types)
1862 .await
1863 {
1864 Ok(edges) => edges,
1865 Err(e) => {
1866 tracing::warn!(error = %e, "recall_graph_view: zoom_out bfs failed");
1867 continue;
1868 }
1869 };
1870
1871 let mut added = 0usize;
1872 for n_edge in neighbors {
1873 if added >= neighbor_cap || total_neighbors >= total_neighbor_cap {
1874 break;
1875 }
1876 let key = make_key(&n_edge.fact);
1877 if seen.insert(key) {
1878 rf.neighbors.push(n_edge.fact);
1879 added += 1;
1880 total_neighbors += 1;
1881 }
1882 }
1883 }
1884 }
1885
1886 #[cfg(feature = "profiling")]
1887 tracing::Span::current().record("result_count", recalled.len());
1888 Ok(recalled)
1889 }
1890
1891 pub async fn recall_graph_astar(
1899 &self,
1900 query: &str,
1901 limit: usize,
1902 max_hops: u32,
1903 temporal_decay_rate: f64,
1904 edge_types: &[crate::graph::EdgeType],
1905 ) -> Result<Vec<crate::graph::types::GraphFact>, MemoryError> {
1906 let Some(store) = &self.graph_store else {
1907 return Ok(Vec::new());
1908 };
1909 crate::graph::retrieval_astar::graph_recall_astar(
1910 store,
1911 self.qdrant.as_deref(),
1912 &self.provider,
1913 query,
1914 limit,
1915 max_hops,
1916 edge_types,
1917 temporal_decay_rate,
1918 self.hebbian_reinforcement.is_enabled(),
1919 self.hebbian_lr,
1920 self.query_sensitive_cost,
1921 self.embed_timeout,
1922 )
1923 .await
1924 }
1925
1926 pub async fn recall_graph_watercircles(
1934 &self,
1935 query: &str,
1936 limit: usize,
1937 max_hops: u32,
1938 ring_limit: usize,
1939 temporal_decay_rate: f64,
1940 edge_types: &[crate::graph::EdgeType],
1941 ) -> Result<Vec<crate::graph::types::GraphFact>, MemoryError> {
1942 let Some(store) = &self.graph_store else {
1943 return Ok(Vec::new());
1944 };
1945 crate::graph::retrieval_watercircles::graph_recall_watercircles(
1946 store,
1947 self.qdrant.as_deref(),
1948 &self.provider,
1949 query,
1950 limit,
1951 max_hops,
1952 ring_limit,
1953 edge_types,
1954 temporal_decay_rate,
1955 self.hebbian_reinforcement.is_enabled(),
1956 self.hebbian_lr,
1957 self.embed_timeout,
1958 )
1959 .await
1960 }
1961
1962 pub async fn recall_graph_beam(
1970 &self,
1971 query: &str,
1972 limit: usize,
1973 beam_width: usize,
1974 max_hops: u32,
1975 temporal_decay_rate: f64,
1976 edge_types: &[crate::graph::EdgeType],
1977 ) -> Result<Vec<crate::graph::types::GraphFact>, MemoryError> {
1978 let Some(store) = &self.graph_store else {
1979 return Ok(Vec::new());
1980 };
1981 crate::graph::retrieval_beam::graph_recall_beam(
1982 store,
1983 self.qdrant.as_deref(),
1984 &self.provider,
1985 query,
1986 limit,
1987 beam_width,
1988 max_hops,
1989 edge_types,
1990 temporal_decay_rate,
1991 self.hebbian_reinforcement.is_enabled(),
1992 self.hebbian_lr,
1993 self.embed_timeout,
1994 )
1995 .await
1996 }
1997
1998 pub async fn classify_graph_strategy(&self, query: &str) -> String {
2003 crate::graph::strategy_classifier::classify_retrieval_strategy(&self.provider, query).await
2004 }
2005
2006 #[cfg_attr(
2020 feature = "profiling",
2021 tracing::instrument(
2022 name = "memory.recall_graph_hela",
2023 skip_all,
2024 fields(result_count = tracing::field::Empty)
2025 )
2026 )]
2027 pub async fn recall_graph_hela(
2028 &self,
2029 query: &str,
2030 limit: usize,
2031 params: crate::graph::HelaSpreadParams,
2032 ) -> Result<Vec<crate::graph::HelaFact>, MemoryError> {
2033 let Some(store) = &self.graph_store else {
2034 return Ok(Vec::new());
2035 };
2036 let Some(embeddings) = &self.qdrant else {
2037 return Ok(Vec::new());
2038 };
2039
2040 let store = Arc::clone(store);
2041 let embeddings = Arc::clone(embeddings);
2042 let provider = self.provider.clone();
2043 let hebbian_enabled = self.hebbian_reinforcement.is_enabled();
2044 let hebbian_lr = self.hebbian_lr;
2045
2046 let outer_timeout = hela_outer_timeout(¶ms);
2051
2052 let results = tokio::time::timeout(
2053 outer_timeout,
2054 crate::graph::hela_spreading_recall(
2055 &store,
2056 &embeddings,
2057 &provider,
2058 query,
2059 limit,
2060 ¶ms,
2061 hebbian_enabled,
2062 hebbian_lr,
2063 ),
2064 )
2065 .await
2066 .unwrap_or_else(|_| {
2067 tracing::warn!(
2068 outer_timeout_ms = outer_timeout.as_millis(),
2069 "memory.recall_graph_hela: outer timeout exceeded"
2070 );
2071 Ok(Vec::new())
2072 })?;
2073
2074 #[cfg(feature = "profiling")]
2075 tracing::Span::current().record("result_count", results.len());
2076
2077 Ok(results)
2078 }
2079
2080 async fn batch_increment_access_count(
2088 &self,
2089 message_ids: Vec<MessageId>,
2090 ) -> Result<(), MemoryError> {
2091 if message_ids.is_empty() {
2092 return Ok(());
2093 }
2094 self.sqlite.increment_access_counts(&message_ids).await
2095 }
2096
2097 pub async fn has_embedding(&self, message_id: MessageId) -> Result<bool, MemoryError> {
2103 match &self.qdrant {
2104 Some(qdrant) => qdrant.has_embedding(message_id).await,
2105 None => Ok(false),
2106 }
2107 }
2108
2109 pub async fn embed_missing(
2125 &self,
2126 progress_tx: Option<tokio::sync::watch::Sender<Option<super::BackfillProgress>>>,
2127 ) -> Result<usize, MemoryError> {
2128 if self.qdrant.is_none() || !self.effective_embed_provider().supports_embeddings() {
2129 return Ok(0);
2130 }
2131
2132 let total = self.sqlite.count_unembedded_messages().await?;
2133 if total == 0 {
2134 return Ok(0);
2135 }
2136
2137 if let Some(tx) = &progress_tx {
2138 let _ = tx.send(Some(super::BackfillProgress { done: 0, total }));
2139 }
2140
2141 let mut done = 0usize;
2142 let mut succeeded = 0usize;
2143
2144 loop {
2145 const BATCH_SIZE: usize = 32;
2146 const BATCH_SIZE_I64: i64 = 32;
2147 let rows: Vec<_> = self
2148 .sqlite
2149 .stream_unembedded_messages(BATCH_SIZE_I64)
2150 .try_collect()
2151 .await?;
2152
2153 if rows.is_empty() {
2154 break;
2155 }
2156
2157 let batch_len = rows.len();
2158
2159 let results: Vec<bool> = futures::stream::iter(rows)
2160 .map(|(msg_id, conv_id, role, content)| async move {
2161 self.embed_and_store_regular(msg_id, conv_id, &role, &content)
2162 })
2163 .buffer_unordered(4)
2164 .collect()
2165 .await;
2166
2167 for ok in &results {
2168 done += 1;
2169 if *ok {
2170 succeeded += 1;
2171 }
2172 if let Some(tx) = &progress_tx {
2173 let _ = tx.send(Some(super::BackfillProgress { done, total }));
2174 }
2175 }
2176
2177 let batch_succeeded = results.iter().filter(|&&b| b).count();
2178 if batch_succeeded > 0 {
2179 tracing::debug!("Backfill batch: {batch_succeeded}/{batch_len} embedded");
2180 }
2181
2182 if batch_len < BATCH_SIZE {
2183 break;
2184 }
2185 }
2186
2187 if let Some(tx) = &progress_tx {
2188 let _ = tx.send(None);
2189 }
2190
2191 if done > 0 {
2192 tracing::info!("Embedded {succeeded}/{total} missing messages");
2193 }
2194 Ok(succeeded)
2195 }
2196}
2197
2198#[cfg(test)]
2199mod tests {
2200 use super::*;
2201
2202 #[test]
2207 fn hela_outer_timeout_scales_with_spread_depth() {
2208 let step_budget = std::time::Duration::from_millis(80);
2209 let embed_timeout = std::time::Duration::from_secs(5);
2210 let margin = std::time::Duration::from_millis(250);
2211
2212 for depth in 1..=6u32 {
2213 let params = crate::graph::HelaSpreadParams {
2214 spread_depth: depth,
2215 step_budget: Some(step_budget),
2216 embed_timeout: Some(embed_timeout),
2217 ..Default::default()
2218 };
2219 let expected = embed_timeout + step_budget * (depth + 2) + margin;
2220 assert_eq!(
2221 hela_outer_timeout(¶ms),
2222 expected,
2223 "outer timeout must scale with spread_depth={depth}"
2224 );
2225 }
2226 }
2227
2228 #[test]
2232 fn hela_outer_timeout_clamps_spread_depth_above_six() {
2233 let step_budget = std::time::Duration::from_millis(80);
2234 let embed_timeout = std::time::Duration::from_secs(5);
2235 let params_over = crate::graph::HelaSpreadParams {
2236 spread_depth: 50,
2237 step_budget: Some(step_budget),
2238 embed_timeout: Some(embed_timeout),
2239 ..Default::default()
2240 };
2241 let params_clamped = crate::graph::HelaSpreadParams {
2242 spread_depth: 6,
2243 step_budget: Some(step_budget),
2244 embed_timeout: Some(embed_timeout),
2245 ..Default::default()
2246 };
2247 assert_eq!(
2248 hela_outer_timeout(¶ms_over),
2249 hela_outer_timeout(¶ms_clamped),
2250 "spread_depth above 6 must clamp identically to the algorithm's own [1, 6] bound"
2251 );
2252 }
2253
2254 #[test]
2258 fn hela_outer_timeout_falls_back_when_params_disabled() {
2259 let params = crate::graph::HelaSpreadParams {
2260 spread_depth: 2,
2261 step_budget: None,
2262 embed_timeout: None,
2263 ..Default::default()
2264 };
2265 let expected = std::time::Duration::from_secs(5)
2266 + std::time::Duration::from_millis(80) * 4
2267 + std::time::Duration::from_millis(250);
2268 assert_eq!(hela_outer_timeout(¶ms), expected);
2269 }
2270
2271 #[test]
2272 fn embed_context_default_all_none() {
2273 let ctx = EmbedContext::default();
2274 assert!(ctx.tool_name.is_none());
2275 assert!(ctx.exit_code.is_none());
2276 assert!(ctx.timestamp.is_none());
2277 }
2278
2279 #[test]
2280 fn embed_context_fields_set_correctly() {
2281 let ctx = EmbedContext {
2282 tool_name: Some("shell".to_string()),
2283 exit_code: Some(0),
2284 timestamp: Some("2026-04-04T00:00:00Z".to_string()),
2285 };
2286 assert_eq!(ctx.tool_name.as_deref(), Some("shell"));
2287 assert_eq!(ctx.exit_code, Some(0));
2288 assert_eq!(ctx.timestamp.as_deref(), Some("2026-04-04T00:00:00Z"));
2289 }
2290
2291 #[test]
2292 fn embed_context_non_zero_exit_code() {
2293 let ctx = EmbedContext {
2294 tool_name: Some("shell".to_string()),
2295 exit_code: Some(1),
2296 timestamp: None,
2297 };
2298 assert_eq!(ctx.exit_code, Some(1));
2299 assert!(ctx.timestamp.is_none());
2300 }
2301
2302 async fn make_semantic_memory() -> crate::semantic::SemanticMemory {
2303 let sqlite = crate::store::SqliteStore::new(":memory:").await.unwrap();
2304 make_semantic_memory_with_sqlite(sqlite)
2305 }
2306
2307 fn make_semantic_memory_with_sqlite(
2315 sqlite: crate::store::SqliteStore,
2316 ) -> crate::semantic::SemanticMemory {
2317 use std::sync::Arc;
2318 use std::sync::atomic::AtomicU64;
2319 use zeph_llm::any::AnyProvider;
2320 use zeph_llm::mock::MockProvider;
2321
2322 let provider = AnyProvider::Mock(MockProvider::default());
2323 crate::semantic::SemanticMemory {
2324 sqlite,
2325 qdrant: None,
2326 provider,
2327 embed_provider: None,
2328 embedding_model: "test-model".into(),
2329 vector_weight: 0.7,
2330 keyword_weight: 0.3,
2331 temporal_decay: crate::semantic::TemporalDecay::Disabled,
2332 temporal_decay_half_life_days: 30,
2333 mmr_reranking: crate::semantic::MmrReranking::Disabled,
2334 mmr_lambda: 0.7,
2335 importance_scoring: crate::semantic::ImportanceScoring::Disabled,
2336 importance_weight: 0.15,
2337 token_counter: Arc::new(crate::token_counter::TokenCounter::new()),
2338 graph_store: None,
2339 experience: None,
2340 community_detection_failures: Arc::new(AtomicU64::new(0)),
2341 graph_extraction_count: Arc::new(AtomicU64::new(0)),
2342 graph_extraction_failures: Arc::new(AtomicU64::new(0)),
2343 last_qdrant_warn: Arc::new(AtomicU64::new(0)),
2344 tier_boost_semantic: 1.3,
2345 admission_control: None,
2346 quality_gate: None,
2347 key_facts_dedup_threshold: 0.95,
2348 embed_tasks: std::sync::Mutex::new(tokio::task::JoinSet::new()),
2349 retrieval_depth: 0,
2350 search_prompt_template: String::new(),
2351 depth_below_limit_warned: Arc::new(std::sync::atomic::AtomicBool::new(false)),
2352 missing_placeholder_warned: Arc::new(std::sync::atomic::AtomicBool::new(false)),
2353 reasoning: None,
2354 query_bias_correction: crate::semantic::QueryBiasCorrection::Disabled,
2355 query_bias_profile_weight: 0.25,
2356 profile_centroid: tokio::sync::RwLock::new(None),
2357 profile_centroid_ttl_secs: 300,
2358 hebbian_reinforcement: crate::semantic::HebbianReinforcement::Disabled,
2359 hebbian_lr: 0.1,
2360 hebbian_spread: crate::HelaSpreadRuntime::default(),
2361 retrieval_failure_logger: None,
2362 summarization_llm_timeout_secs: 60,
2363 query_sensitive_cost: false,
2364 five_signal: None,
2365 embed_timeout: std::time::Duration::from_secs(5),
2366 graph_cancel: std::sync::Mutex::new(Vec::new()),
2367 }
2368 }
2369
2370 #[tokio::test]
2371 async fn spawn_embed_bg_returns_true_when_capacity_available() {
2372 let memory = make_semantic_memory().await;
2373 let dispatched = memory.spawn_embed_bg(std::future::ready(()));
2374 assert!(
2375 dispatched,
2376 "spawn_embed_bg must return true when a task was successfully spawned"
2377 );
2378 }
2379
2380 #[tokio::test]
2381 async fn spawn_embed_bg_returns_false_at_capacity() {
2382 let memory = make_semantic_memory().await;
2383
2384 {
2386 let mut tasks = memory.embed_tasks.lock().unwrap();
2387 for _ in 0..MAX_EMBED_BG_TASKS {
2388 tasks.spawn(std::future::pending::<()>());
2389 }
2390 }
2391
2392 let dispatched = memory.spawn_embed_bg(std::future::ready(()));
2393 assert!(
2394 !dispatched,
2395 "spawn_embed_bg must return false when the task limit is reached"
2396 );
2397 }
2398
2399 #[test]
2400 fn qdrant_warn_rate_limit_suppresses_within_window() {
2401 assert!(
2403 should_emit_qdrant_warn(0, 100, 10),
2404 "first call must not be suppressed"
2405 );
2406
2407 assert!(
2409 !should_emit_qdrant_warn(100, 105, 10),
2410 "call within 10s window must be suppressed"
2411 );
2412
2413 assert!(
2415 should_emit_qdrant_warn(100, 110, 10),
2416 "call after window expiry must not be suppressed"
2417 );
2418 }
2419
2420 #[cfg(feature = "test-utils")]
2431 #[tokio::test]
2432 #[ignore = "requires Docker"]
2433 async fn apply_five_signal_scoring_decodes_created_at_on_postgres() {
2434 use std::sync::Arc;
2435 use testcontainers::runners::AsyncRunner as _;
2436 use testcontainers_modules::postgres::Postgres;
2437 use zeph_config::memory::FiveSignalConfig;
2438
2439 let image = Postgres::default();
2440 let container = image.start().await.expect("docker must be available");
2441 let host = container.get_host().await.unwrap();
2442 let port = container.get_host_port_ipv4(5432).await.unwrap();
2443 let url = format!("postgres://postgres:postgres@{host}:{port}/postgres");
2444 let pool = zeph_db::DbConfig { url, pool_size: 5 }
2445 .connect()
2446 .await
2447 .expect("failed to connect to PG");
2448
2449 let pg_store = crate::store::SqliteStore::from_pool(pool.clone())
2450 .await
2451 .unwrap();
2452 let cid = pg_store.create_conversation().await.unwrap();
2453
2454 let session_start = 1_700_000_000_i64;
2459 let fresh = pg_store.save_message(cid, "user", "fresh").await.unwrap();
2460 let stale = pg_store.save_message(cid, "user", "stale").await.unwrap();
2461
2462 for (id, epoch) in [(fresh, session_start), (stale, session_start + 20 * 86_400)] {
2463 #[expect(clippy::cast_precision_loss)]
2464 let epoch_f = epoch as f64;
2465 sqlx::query(zeph_db::sql!(
2466 "UPDATE messages SET created_at = to_timestamp(?) WHERE id = ?"
2467 ))
2468 .bind(epoch_f)
2469 .bind(id)
2470 .execute(&pool)
2471 .await
2472 .unwrap();
2473 }
2474
2475 let graph_store = Arc::new(crate::graph::GraphStore::new(pool.clone()));
2476 let config = FiveSignalConfig {
2477 w_recency: 0.0,
2478 w_relevance: 0.0,
2479 w_frequency: 0.0,
2480 w_causal: 0.0,
2481 w_novelty: 1.0,
2482 ..FiveSignalConfig::default()
2483 };
2484 let fs = crate::five_signal::FiveSignalRuntime::new(
2485 config,
2486 pool,
2487 graph_store,
2488 None,
2489 session_start,
2490 "sess-novelty-test",
2491 );
2492
2493 let memory = make_semantic_memory_with_sqlite(pg_store);
2499 let mut ranked = vec![(fresh, 0.0), (stale, 0.0)];
2500 memory
2501 .apply_five_signal_scoring(&mut ranked, &fs, None)
2502 .await;
2503
2504 assert_eq!(
2505 ranked[0].0, fresh,
2506 "fresher message (created_at == session_start) must rank first by novelty"
2507 );
2508 assert!(
2509 (ranked[0].1 - 1.0).abs() < 1e-6,
2510 "message created at session_start must have novelty ~1.0, got {}",
2511 ranked[0].1
2512 );
2513 assert!(
2514 ranked[1].1 < ranked[0].1,
2515 "message created 20 days later must have strictly lower novelty"
2516 );
2517 }
2518
2519 #[test]
2520 fn qdrant_warn_rate_limit_shared_across_concurrent_sites() {
2521 let shared = Arc::new(AtomicU64::new(0));
2526
2527 let site_a = Arc::clone(&shared);
2528 let site_b = Arc::clone(&shared);
2529
2530 let now_a = 100u64;
2531 let last_a = site_a.load(Ordering::Relaxed);
2532 if should_emit_qdrant_warn(last_a, now_a, QDRANT_WARN_WINDOW_SECS) {
2533 site_a.store(now_a, Ordering::Relaxed);
2534 }
2535
2536 let now_b = 105u64;
2537 let last_b = site_b.load(Ordering::Relaxed);
2538 let warn_b = should_emit_qdrant_warn(last_b, now_b, QDRANT_WARN_WINDOW_SECS);
2539 assert!(
2540 !warn_b,
2541 "site B must be suppressed because site A already warned within the window"
2542 );
2543 }
2544
2545 #[test]
2546 fn warn_qdrant_ensure_failure_updates_shared_atomic_once() {
2547 let shared = Arc::new(AtomicU64::new(0));
2548 let err = MemoryError::InvalidInput("boom".into());
2549
2550 warn_qdrant_ensure_failure(&shared, "site A", &err);
2551 let after_first = shared.load(Ordering::Relaxed);
2552 assert!(after_first > 0, "first call must record a warn timestamp");
2553
2554 warn_qdrant_ensure_failure(&shared, "site B", &err);
2557 assert_eq!(
2558 shared.load(Ordering::Relaxed),
2559 after_first,
2560 "suppressed call must not overwrite the shared warn timestamp"
2561 );
2562 }
2563}