1use std::collections::{HashMap, HashSet};
4
5use lattice_embed::{EmbeddingModel, MAX_TEXT_BYTES};
6use uuid::Uuid;
7
8use crate::config::{parse_embedding_model_alias, sanitize_key};
9use crate::curation::note_fts_document;
10use crate::embedder_registry::with_embedding_admission;
11use crate::error::{RuntimeError, RuntimeResult};
12use crate::runtime::{KhiveRuntime, NamespaceToken};
13use khive_retrieval::hybrid::{combine_leg_first_appearance, fuse_labelled, HitLabel};
14use khive_score::DeterministicScore;
15use khive_storage::types::{
16 PageRequest, TextFilter, TextQueryMode, TextSearchHit, TextSearchRequest, VectorRecord,
17 VectorSearchHit,
18};
19use khive_storage::ContentRef;
20use khive_storage::EntityFilter;
21use khive_types::SubstrateKind;
22
23pub use khive_retrieval::{
24 HybridSearchOutcome, RankScoreKind, SearchHit, SearchSignals, SearchSource,
25};
26
27mod single_engine;
28
29#[cfg(test)]
30mod request_stop_tests;
31
32pub(crate) const EMBEDDING_BATCH_PAGE_SIZE: usize = 256;
34
35#[cfg(any(test, feature = "fault-injection"))]
37std::thread_local! {
38 static BACKFILL_READER_FAIL: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
39}
40
41#[cfg(any(test, feature = "fault-injection"))]
47pub fn arm_backfill_reader_fail() {
48 BACKFILL_READER_FAIL.with(|c| c.set(true));
49}
50
51const RRF_K: usize = 10;
58
59const CANDIDATE_MULTIPLIER: u32 = 4;
61
62pub const EMBEDDING_INPUT_TRUNCATED_WARNING: &str =
64 "embedding input was truncated to the embedder maximum; full content was stored unchanged";
65
66#[derive(Clone, Debug)]
69pub struct DocumentEmbeddingOutcome {
70 pub model_name: String,
71 pub vector: Vec<f32>,
72 pub prepared_text_fingerprint: Option<ContentRef>,
74 pub source_bytes: usize,
75 pub embedded_bytes: usize,
76 pub truncated: bool,
77}
78
79#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
81pub struct EmbeddingTruncationReport {
82 pub truncated: u64,
83 pub discarded_bytes: u64,
84}
85
86impl EmbeddingTruncationReport {
87 pub fn observe(&mut self, outcome: &DocumentEmbeddingOutcome) {
88 if outcome.truncated {
89 self.truncated += 1;
90 self.discarded_bytes +=
91 outcome.source_bytes.saturating_sub(outcome.embedded_bytes) as u64;
92 }
93 }
94
95 #[must_use]
96 pub const fn any_truncated(&self) -> bool {
97 self.truncated > 0
98 }
99
100 pub fn merge(&mut self, other: Self) {
102 self.truncated = self.truncated.saturating_add(other.truncated);
103 self.discarded_bytes = self.discarded_bytes.saturating_add(other.discarded_bytes);
104 }
105}
106
107pub fn document_embedding_budget(model_name: &str) -> usize {
109 parse_embedding_model_alias(model_name)
110 .and_then(|model| model.document_instruction())
111 .map_or(MAX_TEXT_BYTES, |prefix| {
112 MAX_TEXT_BYTES.saturating_sub(prefix.len())
113 })
114}
115
116pub fn bounded_embedding_input(text: &str, max_bytes: usize) -> (&str, bool) {
118 if text.len() <= max_bytes {
119 return (text, false);
120 }
121
122 let end = text
123 .char_indices()
124 .map(|(index, _)| index)
125 .take_while(|index| *index <= max_bytes)
126 .last()
127 .unwrap_or(0);
128 (&text[..end], true)
129}
130
131fn prepared_document_fingerprint(text: &str, model: EmbeddingModel) -> ContentRef {
132 match model.document_instruction() {
133 Some(prefix) => VectorRecord::fingerprint_text(&format!("{prefix}{text}")),
134 None => VectorRecord::fingerprint_text(text),
135 }
136}
137
138impl KhiveRuntime {
139 fn require_default_embedder(&self) -> RuntimeResult<&str> {
140 let model_name = self.default_embedder_name();
141 if model_name.is_empty() {
142 return Err(RuntimeError::Unconfigured("embedding_model".into()));
143 }
144 Ok(model_name)
145 }
146
147 pub async fn embed(&self, text: &str) -> RuntimeResult<Vec<f32>> {
152 let model_name = self.require_default_embedder()?;
153 self.embed_with_model(model_name, text).await
154 }
155
156 pub async fn embed_with_model(&self, model_name: &str, text: &str) -> RuntimeResult<Vec<f32>> {
172 let model = parse_embedding_model_alias(model_name);
173 let service = self.embedder(model_name).await?;
174 let emb_model = model.unwrap_or_default();
175 crate::usage::count(crate::usage::UsageUnit::EmbedCalls, 1);
179 let out = with_embedding_admission(service.embed_one(text, emb_model)).await;
180 out
181 }
182
183 pub async fn embed_document_with_model_outcome(
202 &self,
203 model_name: &str,
204 text: &str,
205 ) -> RuntimeResult<DocumentEmbeddingOutcome> {
206 self.embed_document_with_model_outcome_inner(None, model_name, text)
207 .await
208 }
209
210 pub(crate) async fn embed_document_with_model_outcome_for_token(
211 &self,
212 token: &NamespaceToken,
213 model_name: &str,
214 text: &str,
215 ) -> RuntimeResult<DocumentEmbeddingOutcome> {
216 self.embed_document_with_model_outcome_inner(Some(token), model_name, text)
217 .await
218 }
219
220 async fn embed_document_with_model_outcome_inner(
221 &self,
222 token: Option<&NamespaceToken>,
223 model_name: &str,
224 text: &str,
225 ) -> RuntimeResult<DocumentEmbeddingOutcome> {
226 let model = parse_embedding_model_alias(model_name);
227 let (service, audited_document_preparation) = self
228 .embedder_with_input_attestation(model_name, token)
229 .await?;
230 let emb_model = model.unwrap_or_default();
231 let source_bytes = text.len();
232 let (text, truncated) =
233 bounded_embedding_input(text, document_embedding_budget(model_name));
234 let embedded_bytes = text.len();
235 crate::usage::count(crate::usage::UsageUnit::EmbedCalls, 1);
237 let embeddings =
238 with_embedding_admission(service.embed_passage(&[text.to_string()], emb_model)).await;
239 let mut vectors = embeddings?;
240 if vectors.len() != 1 {
241 return Err(RuntimeError::Internal(format!(
242 "embed_passage returned {} vectors for 1 input",
243 vectors.len()
244 )));
245 }
246 let out = vectors.pop().expect("checked len == 1 above");
247 Ok(DocumentEmbeddingOutcome {
248 model_name: model_name.to_owned(),
249 vector: out,
250 prepared_text_fingerprint: audited_document_preparation
251 .then(|| prepared_document_fingerprint(text, emb_model)),
252 source_bytes,
253 embedded_bytes,
254 truncated,
255 })
256 }
257
258 pub async fn embed_query_with_model(
270 &self,
271 model_name: &str,
272 text: &str,
273 ) -> RuntimeResult<Vec<f32>> {
274 self.embed_query_with_model_inner(None, model_name, text)
275 .await
276 }
277
278 pub(crate) async fn embed_query_with_model_for_token(
279 &self,
280 token: &NamespaceToken,
281 model_name: &str,
282 text: &str,
283 ) -> RuntimeResult<Vec<f32>> {
284 self.embed_query_with_model_inner(Some(token), model_name, text)
285 .await
286 }
287
288 async fn embed_query_with_model_inner(
289 &self,
290 token: Option<&NamespaceToken>,
291 model_name: &str,
292 text: &str,
293 ) -> RuntimeResult<Vec<f32>> {
294 let model = parse_embedding_model_alias(model_name);
295 let service = match token {
296 Some(token) => self.embedder_with_token(token, model_name).await?,
297 None => self.embedder(model_name).await?,
298 };
299 let texts = [text.to_string()];
300 let emb_model = model.unwrap_or_default();
301 crate::usage::count(crate::usage::UsageUnit::EmbedCalls, 1);
303 let embeddings = match emb_model {
304 EmbeddingModel::BgeSmallEnV15
305 | EmbeddingModel::BgeBaseEnV15
306 | EmbeddingModel::BgeLargeEnV15 => {
307 with_embedding_admission(service.embed(&texts, emb_model)).await
308 }
309 _ => with_embedding_admission(service.embed_query(&texts, emb_model)).await,
310 };
311 let out = embeddings?
312 .into_iter()
313 .next()
314 .ok_or_else(|| RuntimeError::Internal("embed_query returned empty vec".into()))?;
315 Ok(out)
316 }
317
318 pub async fn embed_document(&self, text: &str) -> RuntimeResult<Vec<f32>> {
327 let outcome = self.embed_document_outcome(text).await?;
328 if outcome.truncated {
329 return Err(RuntimeError::InvalidInput(format!(
330 "embedding input truncated from {} to {} bytes; use embed_document_outcome to inspect the bounded vector",
331 outcome.source_bytes, outcome.embedded_bytes
332 )));
333 }
334 Ok(outcome.vector)
335 }
336
337 pub async fn embed_document_outcome(
339 &self,
340 text: &str,
341 ) -> RuntimeResult<DocumentEmbeddingOutcome> {
342 let model_name = self.require_default_embedder()?;
343 self.embed_document_with_model_outcome(model_name, text)
344 .await
345 }
346
347 pub async fn embed_query(&self, text: &str) -> RuntimeResult<Vec<f32>> {
354 let model_name = self.require_default_embedder()?;
355 self.embed_query_with_model(model_name, text).await
356 }
357
358 async fn embed_query_for_token(
359 &self,
360 token: &NamespaceToken,
361 text: &str,
362 ) -> RuntimeResult<Vec<f32>> {
363 let model_name = self.require_default_embedder()?;
364 self.embed_query_with_model_for_token(token, model_name, text)
365 .await
366 }
367
368 pub async fn embed_batch(&self, texts: &[String]) -> RuntimeResult<Vec<Vec<f32>>> {
376 if texts.is_empty() {
377 return Ok(vec![]);
378 }
379 let model_name = self.require_default_embedder()?;
380 self.embed_batch_with_model(model_name, texts).await
381 }
382
383 pub async fn embed_batch_with_model(
388 &self,
389 model_name: &str,
390 texts: &[String],
391 ) -> RuntimeResult<Vec<Vec<f32>>> {
392 if texts.is_empty() {
393 return Ok(vec![]);
394 }
395 let model = parse_embedding_model_alias(model_name);
396 let service = self.embedder(model_name).await?;
397 let emb_model = model.unwrap_or_default();
398 let out = with_embedding_admission(service.embed(texts, emb_model)).await;
399 crate::usage::count(crate::usage::UsageUnit::EmbedCalls, texts.len() as u64);
400 out
401 }
402
403 pub async fn embed_document_batch_with_model(
417 &self,
418 model_name: &str,
419 texts: &[String],
420 ) -> RuntimeResult<Vec<Vec<f32>>> {
421 if texts.is_empty() {
422 return Ok(vec![]);
423 }
424 let outcomes = self
425 .embed_document_batch_with_model_outcomes(model_name, texts)
426 .await?;
427 let mut report = EmbeddingTruncationReport::default();
428 for outcome in &outcomes {
429 report.observe(outcome);
430 }
431 if report.any_truncated() {
432 return Err(RuntimeError::InvalidInput(format!(
433 "embedding input truncated for {} documents ({} discarded bytes); use embed_document_batch_with_model_outcomes to inspect the bounded vectors",
434 report.truncated, report.discarded_bytes
435 )));
436 }
437 Ok(outcomes.into_iter().map(|outcome| outcome.vector).collect())
438 }
439
440 pub async fn embed_document_batch_with_model_outcomes(
441 &self,
442 model_name: &str,
443 texts: &[String],
444 ) -> RuntimeResult<Vec<DocumentEmbeddingOutcome>> {
445 self.embed_document_batch_with_model_outcomes_inner(None, model_name, texts)
446 .await
447 }
448
449 pub(crate) async fn embed_document_batch_with_model_outcomes_for_token(
450 &self,
451 token: &NamespaceToken,
452 model_name: &str,
453 texts: &[String],
454 ) -> RuntimeResult<Vec<DocumentEmbeddingOutcome>> {
455 self.embed_document_batch_with_model_outcomes_inner(Some(token), model_name, texts)
456 .await
457 }
458
459 async fn embed_document_batch_with_model_outcomes_inner(
460 &self,
461 token: Option<&NamespaceToken>,
462 model_name: &str,
463 texts: &[String],
464 ) -> RuntimeResult<Vec<DocumentEmbeddingOutcome>> {
465 if texts.is_empty() {
466 return Ok(vec![]);
467 }
468 let model = parse_embedding_model_alias(model_name);
469 let (service, audited_document_preparation) = self
470 .embedder_with_input_attestation(model_name, token)
471 .await?;
472 let emb_model = model.unwrap_or_default();
473 let budget = document_embedding_budget(model_name);
474 if token.is_some() {
475 crate::usage::count(crate::usage::UsageUnit::EmbedCalls, texts.len() as u64);
477 }
478 let out = if texts.iter().all(|text| text.len() <= budget) {
479 with_embedding_admission(service.embed_passage(texts, emb_model)).await
480 } else {
481 let bounded_texts: Vec<String> = texts
482 .iter()
483 .map(|text| bounded_embedding_input(text, budget).0.to_owned())
484 .collect();
485 with_embedding_admission(service.embed_passage(&bounded_texts, emb_model)).await
486 };
487 if token.is_none() {
488 crate::usage::count(crate::usage::UsageUnit::EmbedCalls, texts.len() as u64);
489 }
490 let vectors = out?;
491 if vectors.len() != texts.len() {
492 return Err(RuntimeError::Internal(format!(
493 "embed_passage returned {} vectors for {} inputs",
494 vectors.len(),
495 texts.len()
496 )));
497 }
498 Ok(texts
499 .iter()
500 .zip(vectors)
501 .map(|(text, vector)| {
502 let (bounded, truncated) = bounded_embedding_input(text, budget);
503 DocumentEmbeddingOutcome {
504 model_name: model_name.to_owned(),
505 vector,
506 prepared_text_fingerprint: audited_document_preparation
507 .then(|| prepared_document_fingerprint(bounded, emb_model)),
508 source_bytes: text.len(),
509 embedded_bytes: bounded.len(),
510 truncated,
511 }
512 })
513 .collect())
514 }
515
516 pub async fn embed_document_batch(&self, texts: &[String]) -> RuntimeResult<Vec<Vec<f32>>> {
525 if texts.is_empty() {
526 return Ok(vec![]);
527 }
528 let model_name = self.require_default_embedder()?;
529 self.embed_document_batch_with_model(model_name, texts)
530 .await
531 }
532
533 pub async fn embed_document_batch_outcomes(
535 &self,
536 texts: &[String],
537 ) -> RuntimeResult<Vec<DocumentEmbeddingOutcome>> {
538 if texts.is_empty() {
539 return Ok(vec![]);
540 }
541 let model_name = self.require_default_embedder()?;
542 self.embed_document_batch_with_model_outcomes(model_name, texts)
543 .await
544 }
545
546 pub async fn embed_query_batch_with_model(
553 &self,
554 model_name: &str,
555 texts: &[String],
556 ) -> RuntimeResult<Vec<Vec<f32>>> {
557 if texts.is_empty() {
558 return Ok(vec![]);
559 }
560 let model = parse_embedding_model_alias(model_name);
561 let service = self.embedder(model_name).await?;
562 let emb_model = model.unwrap_or_default();
563 let out = match emb_model {
564 EmbeddingModel::BgeSmallEnV15
565 | EmbeddingModel::BgeBaseEnV15
566 | EmbeddingModel::BgeLargeEnV15 => {
567 with_embedding_admission(service.embed(texts, emb_model)).await
568 }
569 _ => with_embedding_admission(service.embed_query(texts, emb_model)).await,
570 };
571 crate::usage::count(crate::usage::UsageUnit::EmbedCalls, texts.len() as u64);
572 out
573 }
574
575 pub(crate) async fn note_search_vector_search(
579 &self,
580 token: &NamespaceToken,
581 query_embedding: Option<Vec<f32>>,
582 query_text: &str,
583 top_k: u32,
584 ) -> RuntimeResult<Vec<VectorSearchHit>> {
585 let embedding = match query_embedding {
586 Some(embedding) => embedding,
587 None => self.embed_query_for_token(token, query_text).await?,
588 };
589 let model = self.default_embedder_name();
590 if !model.is_empty() {
591 if let Some(provider) = self.note_search_ann_provider()? {
592 if let Some(hits) = provider.search(token, model, &embedding, top_k).await? {
593 crate::note_search_ann::record_ann_route();
594 crate::usage::count(crate::usage::UsageUnit::VectorPasses, 1);
595 return Ok(hits);
596 }
597 }
598 }
599 crate::note_search_ann::record_fallback_route();
600 self.vector_search(
601 token,
602 Some(embedding),
603 None,
604 top_k,
605 Some(SubstrateKind::Note),
606 )
607 .await
608 }
609
610 #[allow(clippy::too_many_arguments)]
656 pub async fn hybrid_search(
657 &self,
658 token: &NamespaceToken,
659 query_text: &str,
660 query_vector: Option<Vec<f32>>,
661 limit: u32,
662 entity_kind: Option<&str>,
663 entity_type: Option<&str>,
664 tags_any: &[String],
665 properties_filter: Option<&serde_json::Value>,
666 ) -> RuntimeResult<Vec<SearchHit>> {
667 self.hybrid_search_with_text_mode(
668 token,
669 query_text,
670 query_vector,
671 limit,
672 entity_kind,
673 entity_type,
674 tags_any,
675 properties_filter,
676 TextQueryMode::Plain,
677 )
678 .await
679 }
680
681 #[allow(clippy::too_many_arguments)]
683 pub async fn hybrid_search_with_text_mode(
684 &self,
685 token: &NamespaceToken,
686 query_text: &str,
687 query_vector: Option<Vec<f32>>,
688 limit: u32,
689 entity_kind: Option<&str>,
690 entity_type: Option<&str>,
691 tags_any: &[String],
692 properties_filter: Option<&serde_json::Value>,
693 text_mode: TextQueryMode,
694 ) -> RuntimeResult<Vec<SearchHit>> {
695 let (hits, _vector_error) = self
696 .hybrid_search_inner(
697 token,
698 query_text,
699 query_vector,
700 limit,
701 entity_kind,
702 entity_type,
703 tags_any,
704 properties_filter,
705 text_mode,
706 None,
707 false,
708 None,
709 )
710 .await?;
711 Ok(hits)
712 }
713
714 pub async fn hybrid_search_each_kind(
725 &self,
726 token: &NamespaceToken,
727 query_text: &str,
728 query_vector: Option<Vec<f32>>,
729 limit: u32,
730 entity_kinds: &[&str],
731 ) -> RuntimeResult<Vec<Vec<SearchHit>>> {
732 if entity_kinds.is_empty() {
733 return Ok(Vec::new());
734 }
735 let candidates = limit.saturating_mul(CANDIDATE_MULTIPLIER).max(limit);
736 let (vector_hits, _vector_error) = self
737 .hybrid_vector_stage(token, query_text, query_vector, candidates, None, false)
738 .await?;
739 let mut per_kind = Vec::with_capacity(entity_kinds.len());
740 for &kind in entity_kinds {
741 let (hits, _vector_error) = self
742 .hybrid_search_inner(
743 token,
744 query_text,
745 None,
746 limit,
747 Some(kind),
748 None,
749 &[],
750 None,
751 TextQueryMode::Plain,
752 None,
753 false,
754 Some(vector_hits.clone()),
755 )
756 .await?;
757 per_kind.push(hits);
758 }
759 Ok(per_kind)
760 }
761
762 #[allow(clippy::too_many_arguments)]
767 pub(crate) async fn hybrid_search_with_vector_similarity_floor(
768 &self,
769 token: &NamespaceToken,
770 query_text: &str,
771 query_vector: Option<Vec<f32>>,
772 limit: u32,
773 entity_kind: Option<&str>,
774 entity_type: Option<&str>,
775 tags_any: &[String],
776 properties_filter: Option<&serde_json::Value>,
777 vector_similarity_floor: f64,
778 ) -> RuntimeResult<Vec<SearchHit>> {
779 let (hits, _vector_error) = self
780 .hybrid_search_inner(
781 token,
782 query_text,
783 query_vector,
784 limit,
785 entity_kind,
786 entity_type,
787 tags_any,
788 properties_filter,
789 TextQueryMode::Plain,
790 Some(vector_similarity_floor),
791 false,
792 None,
793 )
794 .await?;
795 Ok(hits)
796 }
797
798 #[allow(clippy::too_many_arguments)]
809 pub async fn hybrid_search_outcome(
810 &self,
811 token: &NamespaceToken,
812 query_text: &str,
813 limit: u32,
814 entity_kind: Option<&str>,
815 entity_type: Option<&str>,
816 tags_any: &[String],
817 properties_filter: Option<&serde_json::Value>,
818 ) -> RuntimeResult<HybridSearchOutcome> {
819 self.hybrid_search_outcome_with_text_mode(
820 token,
821 query_text,
822 limit,
823 entity_kind,
824 entity_type,
825 tags_any,
826 properties_filter,
827 TextQueryMode::Plain,
828 )
829 .await
830 }
831
832 #[allow(clippy::too_many_arguments)]
834 pub async fn hybrid_search_outcome_with_text_mode(
835 &self,
836 token: &NamespaceToken,
837 query_text: &str,
838 limit: u32,
839 entity_kind: Option<&str>,
840 entity_type: Option<&str>,
841 tags_any: &[String],
842 properties_filter: Option<&serde_json::Value>,
843 text_mode: TextQueryMode,
844 ) -> RuntimeResult<HybridSearchOutcome> {
845 let (hits, vector_error) = self
846 .hybrid_search_inner(
847 token,
848 query_text,
849 None,
850 limit,
851 entity_kind,
852 entity_type,
853 tags_any,
854 properties_filter,
855 text_mode,
856 None,
857 true,
858 None,
859 )
860 .await?;
861 Ok(HybridSearchOutcome { hits, vector_error })
862 }
863
864 #[allow(clippy::too_many_arguments)]
867 async fn hybrid_search_inner(
868 &self,
869 token: &NamespaceToken,
870 query_text: &str,
871 query_vector: Option<Vec<f32>>,
872 limit: u32,
873 entity_kind: Option<&str>,
874 entity_type: Option<&str>,
875 tags_any: &[String],
876 properties_filter: Option<&serde_json::Value>,
877 text_mode: TextQueryMode,
878 vector_similarity_floor: Option<f64>,
879 tolerate_vector_error: bool,
880 vector_pool: Option<Vec<VectorSearchHit>>,
881 ) -> RuntimeResult<(Vec<SearchHit>, Option<String>)> {
882 khive_storage::await_request_read_phase("hybrid_search", async {
883 let candidates = limit.saturating_mul(CANDIDATE_MULTIPLIER).max(limit);
884
885 let visible_ns: Vec<String> = token
886 .visible_namespaces()
887 .iter()
888 .map(|ns| ns.as_str().to_owned())
889 .collect();
890 let text_store = self.text(token)?;
894 let text_fut = text_store.search(TextSearchRequest {
895 query: query_text.to_string(),
896 mode: text_mode,
897 filter: Some(TextFilter {
898 namespaces: visible_ns.clone(),
899 record_kinds: entity_kind
906 .map(|kind| vec![kind.to_string()])
907 .unwrap_or_default(),
908 ..TextFilter::default()
909 }),
910 top_k: candidates,
911 snippet_chars: 200,
912 });
913 let text_fut = crate::stage_seam::text_stage(text_fut);
914 let vector_fut = async {
916 match vector_pool {
917 Some(pool) => Ok((pool, None)),
918 None => {
919 self.hybrid_vector_stage(
920 token,
921 query_text,
922 query_vector,
923 candidates,
924 vector_similarity_floor,
925 tolerate_vector_error,
926 )
927 .await
928 }
929 }
930 };
931 let (text_search_result, vector_result) = tokio::join!(text_fut, vector_fut);
932 let text_hits = crate::error::fts_text_leg_or_err(
937 text_search_result.map_err(RuntimeError::from),
938 "hybrid_search",
939 query_text,
940 )?;
941 let (vector_hits, vector_error) = vector_result?;
942
943 let fusion_limit = text_hits.len().saturating_add(vector_hits.len());
948 let mut fused = rrf_fuse(text_hits, vector_hits, fusion_limit, query_text);
949
950 if !fused.is_empty() {
953 let candidate_ids: Vec<Uuid> = fused.iter().map(|h| h.entity_id).collect();
954 let alive_page = self
955 .entities(token)?
956 .query_entities(
957 token.namespace().as_str(),
958 EntityFilter {
959 ids: candidate_ids,
960 kinds: entity_kind.map(|k| vec![k.to_string()]).unwrap_or_default(),
961 entity_types: entity_type
962 .map(|t| vec![t.to_string()])
963 .unwrap_or_default(),
964 namespaces: visible_ns,
965 tags_any: tags_any.to_vec(),
966 ..EntityFilter::default()
967 },
968 PageRequest {
969 offset: 0,
970 limit: u32::try_from(fused.len()).unwrap_or(u32::MAX),
971 },
972 )
973 .await?;
974 let mut entity_meta: HashMap<Uuid, (String, Option<String>)> = HashMap::new();
975 let mut alive: HashSet<Uuid> = HashSet::new();
976 for e in alive_page.items {
977 if let Some(pf) = properties_filter {
980 if !properties_match(e.properties.as_ref(), pf) {
981 continue;
982 }
983 }
984 alive.insert(e.id);
985 entity_meta.insert(e.id, (e.name, e.description));
986 }
987
988 fused.retain(|h| alive.contains(&h.entity_id));
989
990 for hit in &mut fused {
992 if let Some((name, description)) = entity_meta.get(&hit.entity_id) {
993 if hit.title.is_none() {
994 hit.title = Some(name.clone());
995 }
996 if hit.snippet.is_none() {
997 hit.snippet = description.clone();
998 }
999 }
1000 }
1001 }
1002
1003 fused.truncate(limit as usize);
1004 Ok((fused, vector_error))
1005 })
1006 .await?
1007 }
1008
1009 async fn hybrid_vector_stage(
1013 &self,
1014 token: &NamespaceToken,
1015 query_text: &str,
1016 query_vector: Option<Vec<f32>>,
1017 candidates: u32,
1018 vector_similarity_floor: Option<f64>,
1019 tolerate_vector_error: bool,
1020 ) -> RuntimeResult<(Vec<VectorSearchHit>, Option<String>)> {
1021 khive_storage::await_request_read_phase("hybrid_search.vector", async {
1022 let mut vector_error: Option<String> = None;
1023 let mut vector_hits =
1024 if query_vector.is_some() || self.config().embedding_model.is_some() {
1025 match self
1026 .vector_search(
1027 token,
1028 query_vector,
1029 Some(query_text),
1030 candidates,
1031 Some(SubstrateKind::Entity),
1032 )
1033 .await
1034 {
1035 Ok(hits) => hits,
1036 Err(e) if tolerate_vector_error => {
1037 vector_error = Some(e.to_string());
1038 Vec::new()
1039 }
1040 Err(e) => return Err(e),
1041 }
1042 } else {
1043 Vec::new()
1044 };
1045 if let Some(cosine_floor) = vector_similarity_floor {
1046 let score_floor = DeterministicScore::from_f64(cosine_floor);
1049 vector_hits.retain(|hit| hit.score >= score_floor);
1050 }
1051 Ok((vector_hits, vector_error))
1052 })
1053 .await?
1054 }
1055
1056 async fn embed_backfill_page(
1057 &self,
1058 token: &NamespaceToken,
1059 model_name: &str,
1060 inputs: &[(Uuid, String)],
1061 ) -> Vec<Option<DocumentEmbeddingOutcome>> {
1062 let texts: Vec<String> = inputs.iter().map(|(_, text)| text.clone()).collect();
1063 match self
1064 .embed_document_batch_with_model_outcomes_for_token(token, model_name, &texts)
1065 .await
1066 {
1067 Ok(outcomes) => outcomes.into_iter().map(Some).collect(),
1068 Err(error) => {
1069 tracing::warn!(
1070 model = %model_name,
1071 error = %error,
1072 "backfill_missing_embeddings: batch embed failed; retrying records individually"
1073 );
1074 let mut outcomes = Vec::with_capacity(inputs.len());
1075 for (id, text) in inputs {
1076 match self
1077 .embed_document_with_model_outcome_for_token(token, model_name, text)
1078 .await
1079 {
1080 Ok(outcome) => outcomes.push(Some(outcome)),
1081 Err(error) => {
1082 tracing::warn!(
1083 id = %id,
1084 model = %model_name,
1085 error = %error,
1086 "backfill_missing_embeddings: record embed failed"
1087 );
1088 outcomes.push(None);
1089 }
1090 }
1091 }
1092 outcomes
1093 }
1094 }
1095 }
1096
1097 pub async fn backfill_missing_embeddings(&self, token: &NamespaceToken) -> RuntimeResult<u64> {
1110 use khive_storage::types::{SqlRow, SqlStatement, SqlValue};
1111
1112 let model_names = self.registered_embedding_model_names();
1113 if model_names.is_empty() {
1114 tracing::debug!(
1115 "backfill_missing_embeddings: no embedding models registered, skipping"
1116 );
1117 return Ok(0);
1118 }
1119
1120 let ns = token.namespace().as_str().to_string();
1121 let mut total_backfilled = 0u64;
1122
1123 for model_name in &model_names {
1124 let mut model_truncation = EmbeddingTruncationReport::default();
1125 match self.vectors_for_model(token, model_name) {
1126 Ok(_) => {}
1127 Err(error) => {
1128 tracing::warn!(model = %model_name, error = %error,
1129 "backfill_missing_embeddings: vector store unavailable");
1130 continue;
1131 }
1132 };
1133 let vec_table = format!("vec_{}", sanitize_key(model_name));
1135
1136 const PAGE_SIZE: usize = EMBEDDING_BATCH_PAGE_SIZE;
1140 let mut entity_total = 0usize;
1141 let mut entity_cursor = String::new();
1142 loop {
1143 let entity_sql = SqlStatement {
1144 sql: format!(
1145 "SELECT id FROM entities \
1146 WHERE namespace = ?1 AND deleted_at IS NULL AND id > ?3 \
1147 AND id NOT IN (\
1148 SELECT subject_id FROM {vec_table} \
1149 WHERE namespace = ?1 AND embedding_model = ?2 \
1150 ) ORDER BY id LIMIT {PAGE_SIZE}"
1151 ),
1152 params: vec![
1153 SqlValue::Text(ns.clone()),
1154 SqlValue::Text(model_name.clone()),
1155 SqlValue::Text(entity_cursor.clone()),
1156 ],
1157 label: Some("backfill_entities".into()),
1158 };
1159
1160 let entity_rows: Vec<SqlRow> = {
1161 let sql = self.sql();
1162 let reader_result = sql.reader().await;
1163 #[cfg(any(test, feature = "fault-injection"))]
1164 let reader_result = if BACKFILL_READER_FAIL.with(|c| c.get()) {
1165 BACKFILL_READER_FAIL.with(|c| c.set(false));
1166 Err(khive_storage::StorageError::Pool {
1167 operation: "reader".into(),
1168 message: "injected failure".into(),
1169 })
1170 } else {
1171 reader_result
1172 };
1173 let mut reader = reader_result.map_err(RuntimeError::Storage)?;
1174 reader
1175 .query_all(entity_sql)
1176 .await
1177 .map_err(RuntimeError::Storage)?
1178 };
1179
1180 let batch_len = entity_rows.len();
1181 entity_total += batch_len;
1182 if batch_len == 0 {
1183 break;
1184 }
1185 entity_cursor = match entity_rows.last().and_then(|row| row.columns.first()) {
1186 Some(column) => match &column.value {
1187 SqlValue::Text(id) => id.clone(),
1188 _ => {
1189 return Err(RuntimeError::Internal(
1190 "backfill entity ID is not text".into(),
1191 ))
1192 }
1193 },
1194 None => {
1195 return Err(RuntimeError::Internal(
1196 "backfill entity page is empty".into(),
1197 ))
1198 }
1199 };
1200
1201 let entity_store = self.entities(token)?;
1202 let mut entities = Vec::with_capacity(batch_len);
1203 let mut inputs = Vec::with_capacity(batch_len);
1204 for row in &entity_rows {
1205 let Some(SqlValue::Text(id)) = row.columns.first().map(|column| &column.value)
1206 else {
1207 continue;
1208 };
1209 let Ok(id) = id.parse::<Uuid>() else { continue };
1210 let Some(entity) = entity_store.get_entity(id).await? else {
1211 continue;
1212 };
1213 if entity.namespace != ns || entity.deleted_at.is_some() {
1214 continue;
1215 }
1216 let text = crate::curation::entity_embedding_text(&entity);
1217 if text.trim().is_empty() {
1218 continue;
1219 }
1220 inputs.push((id, text));
1221 entities.push(entity);
1222 }
1223 let outcomes = self.embed_backfill_page(token, model_name, &inputs).await;
1224 for (entity, outcome) in entities.into_iter().zip(outcomes) {
1225 let Some(outcome) = outcome else { continue };
1226 model_truncation.observe(&outcome);
1227 match self
1228 .publish_entity_vector_revision(token, &entity, model_name, &outcome.vector)
1229 .await
1230 {
1231 Ok(true) => total_backfilled += 1,
1232 Ok(false) => {}
1233 Err(error) => tracing::warn!(
1234 id = %entity.id, model = %model_name, error = %error,
1235 "backfill_missing_embeddings: entity vector insert failed"
1236 ),
1237 }
1238 }
1239
1240 if batch_len < PAGE_SIZE {
1241 break;
1242 }
1243 }
1244
1245 let text_store = self.text_for_notes(token).ok();
1247 let note_store = self.notes(token).ok();
1248 let mut note_total = 0usize;
1249 let mut note_cursor = String::new();
1252 loop {
1253 let note_sql = SqlStatement {
1256 sql: format!(
1257 "SELECT id FROM notes \
1258 WHERE namespace = ?1 AND deleted_at IS NULL AND id > ?3 \
1259 AND id NOT IN (\
1260 SELECT subject_id FROM {vec_table} \
1261 WHERE namespace = ?1 AND embedding_model = ?2 \
1262 ) ORDER BY id LIMIT {PAGE_SIZE}"
1263 ),
1264 params: vec![
1265 SqlValue::Text(ns.clone()),
1266 SqlValue::Text(model_name.clone()),
1267 SqlValue::Text(note_cursor.clone()),
1268 ],
1269 label: Some("backfill_notes".into()),
1270 };
1271
1272 let note_rows: Vec<SqlRow> = {
1273 let sql = self.sql();
1274 let reader_result = sql.reader().await;
1275 #[cfg(any(test, feature = "fault-injection"))]
1276 let reader_result = if BACKFILL_READER_FAIL.with(|c| c.get()) {
1277 BACKFILL_READER_FAIL.with(|c| c.set(false));
1278 Err(khive_storage::StorageError::Pool {
1279 operation: "reader".into(),
1280 message: "injected failure".into(),
1281 })
1282 } else {
1283 reader_result
1284 };
1285 let mut reader = reader_result.map_err(RuntimeError::Storage)?;
1286 reader
1287 .query_all(note_sql)
1288 .await
1289 .map_err(RuntimeError::Storage)?
1290 };
1291
1292 let batch_len = note_rows.len();
1293 note_total += batch_len;
1294 if batch_len == 0 {
1295 break;
1296 }
1297 note_cursor = match note_rows.last().and_then(|row| row.columns.first()) {
1298 Some(column) => match &column.value {
1299 SqlValue::Text(id) => id.clone(),
1300 _ => {
1301 return Err(RuntimeError::Internal(
1302 "backfill note ID is not text".into(),
1303 ))
1304 }
1305 },
1306 None => {
1307 return Err(RuntimeError::Internal("backfill note page is empty".into()))
1308 }
1309 };
1310
1311 let mut notes_for_batch = Vec::with_capacity(batch_len);
1312 let mut inputs = Vec::with_capacity(batch_len);
1313 for row in ¬e_rows {
1314 let id_str = row.columns.first().and_then(|c| {
1315 if let SqlValue::Text(s) = &c.value {
1316 Some(s.clone())
1317 } else {
1318 None
1319 }
1320 });
1321
1322 let Some(id_str) = id_str else {
1323 continue;
1324 };
1325 let Ok(id) = id_str.parse::<Uuid>() else {
1326 continue;
1327 };
1328
1329 let note = match ¬e_store {
1330 Some(store) => match store.get_note(id).await {
1331 Ok(Some(n)) => n,
1332 _ => continue,
1333 },
1334 None => continue,
1335 };
1336
1337 if note.content.trim().is_empty() {
1338 continue;
1339 }
1340
1341 if model_names.first().map(|n| n.as_str()) == Some(model_name.as_str()) {
1344 if let Some(ref ts) = text_store {
1345 if let Err(e) = ts.upsert_document(note_fts_document(¬e)).await {
1346 tracing::warn!(id = %id, error = %e,
1347 "backfill_missing_embeddings: note FTS upsert failed");
1348 }
1349 }
1350 }
1351
1352 if !self
1353 .embedding_models_for_note_kind(¬e.kind)
1354 .contains(model_name)
1355 {
1356 continue;
1357 }
1358
1359 inputs.push((id, note.content.clone()));
1360 notes_for_batch.push(note);
1361 }
1362
1363 let outcomes = self.embed_backfill_page(token, model_name, &inputs).await;
1364 for (note, outcome) in notes_for_batch.into_iter().zip(outcomes) {
1365 let Some(outcome) = outcome else { continue };
1366 model_truncation.observe(&outcome);
1367 match self
1368 .publish_note_vector_revision(token, ¬e, model_name, &outcome.vector)
1369 .await
1370 {
1371 Ok(true) => total_backfilled += 1,
1372 Ok(false) => {}
1373 Err(error) => tracing::warn!(
1374 id = %note.id, model = %model_name, error = %error,
1375 "backfill_missing_embeddings: note vector insert failed"
1376 ),
1377 }
1378 }
1379
1380 if batch_len < PAGE_SIZE {
1381 break;
1382 }
1383 }
1384
1385 tracing::info!(
1386 model = %model_name,
1387 namespace = %ns,
1388 entities = entity_total,
1389 notes = note_total,
1390 truncated = model_truncation.truncated,
1391 discarded_bytes = model_truncation.discarded_bytes,
1392 "backfill_missing_embeddings: model pass complete"
1393 );
1394 }
1395
1396 tracing::info!(
1397 namespace = %ns,
1398 total_backfilled = total_backfilled,
1399 "backfill_missing_embeddings: finished"
1400 );
1401
1402 Ok(total_backfilled)
1403 }
1404
1405 pub async fn sweep_orphan_vectors(
1421 &self,
1422 token: &NamespaceToken,
1423 max_delete_per_model: u32,
1424 dry_run: bool,
1425 ) -> RuntimeResult<u64> {
1426 use khive_storage::types::OrphanSweepConfig;
1427 use khive_storage::StorageError;
1428
1429 let model_names = self.registered_embedding_model_names();
1430 if model_names.is_empty() {
1431 tracing::debug!("sweep_orphan_vectors: no embedding models registered, skipping");
1432 return Ok(0);
1433 }
1434
1435 let ns = token.namespace().as_str().to_string();
1436 let mut total_deleted = 0u64;
1437
1438 for model_name in &model_names {
1439 let store = match self.vectors_for_model(token, model_name) {
1440 Ok(s) => s,
1441 Err(e) => {
1442 tracing::warn!(
1443 model = %model_name,
1444 error = %e,
1445 "sweep_orphan_vectors: failed to get vector store, skipping model"
1446 );
1447 continue;
1448 }
1449 };
1450
1451 let caps = store.capabilities();
1452 if !caps.supports_orphan_sweep {
1453 tracing::debug!(
1454 model = %model_name,
1455 "sweep_orphan_vectors: backend does not support orphan sweep, skipping"
1456 );
1457 continue;
1458 }
1459
1460 let config = OrphanSweepConfig {
1461 subject_id_allowlist: None,
1462 namespaces: vec![ns.clone()],
1463 substrate_kinds: vec![],
1464 max_delete: max_delete_per_model,
1465 dry_run,
1466 };
1467
1468 match store.orphan_sweep(&config).await {
1469 Ok(result) => {
1470 tracing::info!(
1471 model = %model_name,
1472 namespace = %ns,
1473 scanned = result.scanned,
1474 deleted = result.deleted,
1475 would_delete = result.would_delete,
1476 dry_run = dry_run,
1477 "sweep_orphan_vectors: sweep complete"
1478 );
1479 total_deleted += result.deleted;
1480 }
1481 Err(StorageError::Unsupported { .. }) => {
1482 tracing::debug!(
1483 model = %model_name,
1484 "sweep_orphan_vectors: backend returned Unsupported, skipping"
1485 );
1486 }
1487 Err(e) => {
1488 tracing::warn!(
1489 model = %model_name,
1490 error = %e,
1491 "sweep_orphan_vectors: sweep failed, continuing with other models"
1492 );
1493 }
1494 }
1495 }
1496
1497 tracing::info!(
1498 namespace = %ns,
1499 total_deleted = total_deleted,
1500 dry_run = dry_run,
1501 "sweep_orphan_vectors: finished"
1502 );
1503
1504 Ok(total_deleted)
1505 }
1506}
1507
1508pub(crate) fn properties_match(
1513 properties: Option<&serde_json::Value>,
1514 filter: &serde_json::Value,
1515) -> bool {
1516 let required = match filter.as_object() {
1517 Some(obj) if !obj.is_empty() => obj,
1518 _ => return true,
1519 };
1520 let actual = match properties.and_then(serde_json::Value::as_object) {
1521 Some(obj) => obj,
1522 None => return false,
1523 };
1524 required
1525 .iter()
1526 .all(|(k, v)| actual.get(k).is_some_and(|av| av == v))
1527}
1528
1529const EXACT_MATCH_BOOST: f64 = 0.5;
1533
1534fn rrf_fuse(
1543 text_hits: Vec<TextSearchHit>,
1544 vector_hits: Vec<VectorSearchHit>,
1545 limit: usize,
1546 query_text: &str,
1547) -> Vec<SearchHit> {
1548 let mut text_arm = Vec::new();
1549 for (rank, hit) in text_hits.into_iter().enumerate() {
1550 let label = HitLabel {
1551 rank,
1552 signals: SearchSignals {
1553 vector_similarity: None,
1554 keyword_score: Some(hit.score),
1555 },
1556 source: SearchSource::Text,
1557 title: hit.title,
1558 snippet: hit.snippet,
1559 };
1560 text_arm.push((hit.subject_id, label));
1561 }
1562 let mut vector_arm = Vec::new();
1563 for (rank, hit) in vector_hits.into_iter().enumerate() {
1564 let label = HitLabel {
1565 rank,
1566 signals: SearchSignals {
1567 vector_similarity: Some(hit.score),
1568 keyword_score: None,
1569 },
1570 source: SearchSource::Vector,
1571 title: None,
1572 snippet: None,
1573 };
1574 vector_arm.push((hit.subject_id, label));
1575 }
1576
1577 let fused = fuse_labelled(
1578 vec![text_arm, vector_arm],
1579 RRF_K,
1580 combine_leg_first_appearance,
1581 );
1582
1583 let query_lower = query_text.to_lowercase();
1585 let boost = DeterministicScore::from_f64(EXACT_MATCH_BOOST);
1586 let mut hits = Vec::with_capacity(fused.len());
1587 for (entity_id, score, label) in fused {
1588 let title_lower = label.title.as_deref().map(str::to_lowercase);
1589 let exact = title_lower.as_deref() == Some(query_lower.as_str());
1590 let score = if exact { score + boost } else { score };
1591 hits.push(SearchHit {
1592 entity_id,
1593 score,
1594 rank_score_kind: RankScoreKind::Rrf,
1595 signals: label.signals,
1596 source: label.source,
1597 title: label.title,
1598 snippet: label.snippet,
1599 });
1600 }
1601
1602 hits.sort_by(|a, b| {
1603 khive_score::cmp_desc_then_id(a.score, &a.entity_id, b.score, &b.entity_id)
1604 });
1605 hits.truncate(limit);
1606 hits
1607}
1608
1609#[cfg(test)]
1610mod rrf_fuse_label_tests;
1611
1612#[cfg(test)]
1613mod tests {
1614 use super::*;
1615 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
1616 use std::sync::Arc;
1617
1618 use crate::runtime::{KhiveRuntime, NamespaceToken, RuntimeConfig};
1619 use khive_score::rrf_score;
1620 use khive_storage::types::{TextSearchHit, VectorSearchHit};
1621 use khive_types::namespace::Namespace;
1622 use lattice_embed::{EmbedError, EmbeddingModel};
1623
1624 fn break_vector_arm(fail: &AtomicBool) {
1627 fail.store(true, Ordering::SeqCst);
1628 }
1629
1630 struct ConstantEmbeddingService {
1634 dimensions: usize,
1635 fail: Arc<AtomicBool>,
1636 }
1637
1638 #[async_trait::async_trait]
1639 impl EmbeddingService for ConstantEmbeddingService {
1640 async fn embed(
1641 &self,
1642 texts: &[String],
1643 _model: EmbeddingModel,
1644 ) -> Result<Vec<Vec<f32>>, EmbedError> {
1645 if self.fail.load(Ordering::SeqCst) {
1646 return Err(EmbedError::ModelInitialization(
1647 "injected vector-arm failure".to_string(),
1648 ));
1649 }
1650 Ok(texts.iter().map(|_| vec![1.0; self.dimensions]).collect())
1651 }
1652
1653 fn supports_model(&self, _model: EmbeddingModel) -> bool {
1654 true
1655 }
1656
1657 fn name(&self) -> &'static str {
1658 "hybrid-search-test-constant-embedding"
1659 }
1660 }
1661
1662 struct ConstantEmbedderProvider {
1663 name: String,
1664 dimensions: usize,
1665 fail: Arc<AtomicBool>,
1666 }
1667
1668 #[async_trait::async_trait]
1669 impl EmbedderProvider for ConstantEmbedderProvider {
1670 fn name(&self) -> &str {
1671 &self.name
1672 }
1673
1674 fn dimensions(&self) -> usize {
1675 self.dimensions
1676 }
1677
1678 async fn build(&self) -> RuntimeResult<Arc<dyn EmbeddingService>> {
1679 Ok(Arc::new(ConstantEmbeddingService {
1680 dimensions: self.dimensions,
1681 fail: Arc::clone(&self.fail),
1682 }))
1683 }
1684 }
1685
1686 fn runtime_with_constant_embeddings() -> KhiveRuntime {
1689 runtime_with_toggleable_embeddings().0
1690 }
1691
1692 fn runtime_with_toggleable_embeddings() -> (KhiveRuntime, Arc<AtomicBool>) {
1693 let model = EmbeddingModel::AllMiniLmL6V2;
1694 let runtime = KhiveRuntime::new(RuntimeConfig {
1695 db_path: None,
1696 embedding_model: Some(model),
1697 packs: vec!["kg".to_string()],
1698 ..RuntimeConfig::no_embeddings()
1699 })
1700 .expect("in-memory runtime");
1701 let fail = Arc::new(AtomicBool::new(false));
1702 runtime.register_embedder(ConstantEmbedderProvider {
1703 name: model.to_string(),
1704 dimensions: model.dimensions(),
1705 fail: Arc::clone(&fail),
1706 });
1707 (runtime, fail)
1708 }
1709
1710 #[test]
1711 fn bounded_embedding_input_reserves_prefix_and_preserves_utf8() {
1712 assert_eq!(
1713 document_embedding_budget("multilingual-e5-base"),
1714 MAX_TEXT_BYTES - "passage: ".len()
1715 );
1716
1717 let input = format!("{}\u{1f980}tail", "a".repeat(MAX_TEXT_BYTES - 1));
1718 let (bounded, truncated) = bounded_embedding_input(&input, MAX_TEXT_BYTES);
1719 assert!(truncated);
1720 assert_eq!(bounded.len(), MAX_TEXT_BYTES - 1);
1721 assert!(bounded.is_char_boundary(bounded.len()));
1722 assert!(!bounded.contains('\u{1f980}'));
1723 }
1724
1725 #[test]
1726 fn bounded_embedding_input_leaves_normal_text_unchanged() {
1727 let input = "normal byte-identical embedding input";
1728 let (bounded, truncated) = bounded_embedding_input(input, MAX_TEXT_BYTES);
1729 assert!(!truncated);
1730 assert_eq!(bounded, input);
1731 assert_eq!(bounded.as_ptr(), input.as_ptr());
1732 }
1733
1734 fn text_hit(id: Uuid, rank: u32, title: &str) -> TextSearchHit {
1735 TextSearchHit {
1736 subject_id: id,
1737 score: DeterministicScore::from_f64(1.0),
1738 rank,
1739 title: Some(title.to_string()),
1740 snippet: Some("...".to_string()),
1741 }
1742 }
1743
1744 fn vector_hit(id: Uuid, rank: u32) -> VectorSearchHit {
1745 VectorSearchHit {
1746 subject_id: id,
1747 score: DeterministicScore::from_f64(0.9),
1748 rank,
1749 }
1750 }
1751
1752 #[test]
1753 fn rrf_evidence_keeps_absence_distinct_from_measured_zero() {
1754 let text_id = Uuid::from_u128(1);
1755 let vector_id = Uuid::from_u128(2);
1756 let both_id = Uuid::from_u128(3);
1757 let quarter = DeterministicScore::from_raw(1_i64 << 30);
1758 let half = DeterministicScore::from_raw(1_i64 << 31);
1759 let text = vec![
1760 TextSearchHit {
1761 score: DeterministicScore::ZERO,
1762 ..text_hit(text_id, 1, "text")
1763 },
1764 TextSearchHit {
1765 score: quarter,
1766 ..text_hit(both_id, 2, "both")
1767 },
1768 text_hit(both_id, 3, "duplicate"),
1769 ];
1770 let vector = vec![
1771 VectorSearchHit {
1772 score: DeterministicScore::ZERO,
1773 ..vector_hit(vector_id, 1)
1774 },
1775 VectorSearchHit {
1776 score: half,
1777 ..vector_hit(both_id, 2)
1778 },
1779 vector_hit(both_id, 3),
1780 ];
1781 let hits = rrf_fuse(text, vector, 10, "unmatched");
1782 assert_eq!(hits.len(), 3);
1783 for (id, source, signals) in [
1784 (
1785 text_id,
1786 SearchSource::Text,
1787 SearchSignals {
1788 vector_similarity: None,
1789 keyword_score: Some(DeterministicScore::ZERO),
1790 },
1791 ),
1792 (
1793 vector_id,
1794 SearchSource::Vector,
1795 SearchSignals {
1796 vector_similarity: Some(DeterministicScore::ZERO),
1797 keyword_score: None,
1798 },
1799 ),
1800 (
1801 both_id,
1802 SearchSource::Both,
1803 SearchSignals {
1804 vector_similarity: Some(half),
1805 keyword_score: Some(quarter),
1806 },
1807 ),
1808 ] {
1809 let hit = hits.iter().find(|hit| hit.entity_id == id).unwrap();
1810 assert_eq!(hit.rank_score_kind, RankScoreKind::Rrf);
1811 assert_eq!(hit.source, source);
1812 assert_eq!(hit.signals, signals);
1813 }
1814 }
1815
1816 #[test]
1817 fn rrf_evidence_golden_preserves_true_ties_across_permutations() {
1818 let a = Uuid::from_u128(1);
1819 let b = Uuid::from_u128(2);
1820 let expected = vec![
1821 (
1822 a,
1823 748_365_513,
1824 RankScoreKind::Rrf,
1825 SearchSignals {
1826 vector_similarity: Some(DeterministicScore::from_raw(1_i64 << 31)),
1827 keyword_score: Some(DeterministicScore::from_raw(1_i64 << 30)),
1828 },
1829 ),
1830 (
1831 b,
1832 748_365_513,
1833 RankScoreKind::Rrf,
1834 SearchSignals {
1835 vector_similarity: Some(DeterministicScore::from_raw(1_i64 << 32)),
1836 keyword_score: Some(DeterministicScore::from_raw(3_i64 << 30)),
1837 },
1838 ),
1839 ];
1840 for (text_ids, vector_ids) in [([a, b], [b, a]), ([b, a], [a, b])] {
1842 for _ in 0..4 {
1843 let text = text_ids
1844 .into_iter()
1845 .enumerate()
1846 .map(|(rank, id)| TextSearchHit {
1847 score: DeterministicScore::from_raw(
1848 if id == a { 1_i64 } else { 3_i64 } << 30,
1849 ),
1850 ..text_hit(id, rank as u32 + 1, "candidate")
1851 })
1852 .collect();
1853 let vector = vector_ids
1854 .into_iter()
1855 .enumerate()
1856 .map(|(rank, id)| VectorSearchHit {
1857 score: DeterministicScore::from_raw(
1858 if id == a { 1_i64 } else { 2_i64 } << 31,
1859 ),
1860 ..vector_hit(id, rank as u32 + 1)
1861 })
1862 .collect();
1863 let hits = rrf_fuse(text, vector, 10, "unmatched");
1864 assert_eq!(hits.len(), 2);
1865 assert_eq!(hits[0].score, hits[1].score);
1866 assert!(hits.iter().all(|hit| hit.source == SearchSource::Both));
1867 let snapshot: Vec<_> = hits
1868 .iter()
1869 .map(|hit| {
1870 (
1871 hit.entity_id,
1872 hit.score.to_raw(),
1873 hit.rank_score_kind,
1874 hit.signals,
1875 )
1876 })
1877 .collect();
1878 assert_eq!(snapshot, expected);
1879 }
1880 }
1881 }
1882
1883 #[test]
1884 fn rrf_fuse_text_only() {
1885 let a = Uuid::new_v4();
1886 let b = Uuid::new_v4();
1887 let text = vec![text_hit(a, 1, "A"), text_hit(b, 2, "B")];
1888 let hits = rrf_fuse(text, vec![], 10, "query");
1889 assert_eq!(hits.len(), 2);
1890 assert_eq!(hits[0].entity_id, a);
1891 assert_eq!(hits[0].source, SearchSource::Text);
1892 assert_eq!(hits[0].title.as_deref(), Some("A"));
1893 }
1894
1895 #[test]
1896 fn rrf_fuse_vector_only() {
1897 let a = Uuid::new_v4();
1898 let hits = rrf_fuse(vec![], vec![vector_hit(a, 1)], 10, "query");
1899 assert_eq!(hits.len(), 1);
1900 assert_eq!(hits[0].source, SearchSource::Vector);
1901 assert!(hits[0].title.is_none());
1902 }
1903
1904 #[test]
1905 fn rrf_fuse_marks_both_when_in_both_lists() {
1906 let id = Uuid::new_v4();
1907 let text = vec![text_hit(id, 1, "A")];
1908 let vec = vec![vector_hit(id, 1)];
1909 let hits = rrf_fuse(text, vec, 10, "query");
1910 assert_eq!(hits.len(), 1);
1911 assert_eq!(hits[0].source, SearchSource::Both);
1912 }
1913
1914 #[test]
1915 fn rrf_fuse_preserves_unique_leg_scores_exactly() {
1916 let text_only = Uuid::new_v4();
1917 let both = Uuid::new_v4();
1918 let vector_only = Uuid::new_v4();
1919 let text = vec![text_hit(text_only, 1, "A"), text_hit(both, 2, "B")];
1920 let vector = vec![vector_hit(both, 1), vector_hit(vector_only, 2)];
1921
1922 let hits = rrf_fuse(text, vector, 10, "query");
1923 let score_for = |id| {
1924 hits.iter()
1925 .find(|hit| hit.entity_id == id)
1926 .expect("expected fused hit")
1927 .score
1928 };
1929
1930 assert_eq!(score_for(text_only), rrf_score(1, RRF_K));
1931 assert_eq!(score_for(both), rrf_score(2, RRF_K) + rrf_score(1, RRF_K));
1932 assert_eq!(score_for(vector_only), rrf_score(2, RRF_K));
1933 }
1934
1935 #[test]
1936 fn rrf_fuse_counts_duplicate_once_per_leg() {
1937 let id = Uuid::new_v4();
1938 let text = vec![text_hit(id, 1, "A"), text_hit(id, 2, "A duplicate")];
1939 let vector = vec![vector_hit(id, 1), vector_hit(id, 2)];
1940
1941 let hits = rrf_fuse(text, vector, 10, "query");
1942
1943 assert_eq!(hits.len(), 1);
1944 assert_eq!(hits[0].source, SearchSource::Both);
1945 assert_eq!(hits[0].score, rrf_score(1, RRF_K) + rrf_score(1, RRF_K));
1946 }
1947
1948 #[test]
1949 fn rrf_fuse_respects_limit() {
1950 let hits: Vec<TextSearchHit> = (0..20)
1951 .map(|i| text_hit(Uuid::new_v4(), i + 1, "x"))
1952 .collect();
1953 let fused = rrf_fuse(hits, vec![], 5, "query");
1954 assert_eq!(fused.len(), 5);
1955 }
1956
1957 #[test]
1958 fn rrf_fuse_orders_higher_score_first() {
1959 let a = Uuid::new_v4();
1961 let b = Uuid::new_v4();
1962 let text = vec![text_hit(a, 1, "A")];
1963 let vec = vec![vector_hit(a, 1), vector_hit(b, 2)];
1964 let hits = rrf_fuse(text, vec, 10, "query");
1965 assert_eq!(hits[0].entity_id, a);
1966 assert_eq!(hits[0].source, SearchSource::Both);
1967 assert!(hits[0].score > hits[1].score);
1968 }
1969
1970 #[test]
1971 fn rrf_fuse_k10_score_spread_exceeds_threshold() {
1972 let ids: Vec<Uuid> = (0..10).map(|_| Uuid::new_v4()).collect();
1975 let text: Vec<TextSearchHit> = ids
1976 .iter()
1977 .enumerate()
1978 .map(|(i, &id)| text_hit(id, (i + 1) as u32, "x"))
1979 .collect();
1980 let hits = rrf_fuse(text, vec![], 10, "query");
1981 assert_eq!(hits.len(), 10);
1982 let top_score = hits[0].score.to_f64();
1983 let bottom_score = hits[9].score.to_f64();
1984 let spread = top_score - bottom_score;
1985 assert!(
1986 spread >= 0.03,
1987 "score spread {spread:.4} between rank 1 and rank 10 must be ≥ 0.03 (was {spread:.4})"
1988 );
1989 }
1990
1991 #[test]
1992 fn rrf_fuse_exact_match_boost_elevates_score() {
1993 let exact_id = Uuid::new_v4();
1996 let other_id = Uuid::new_v4();
1997 let text = vec![
1999 text_hit(other_id, 1, "something else"),
2000 text_hit(exact_id, 2, "FlashAttention"),
2001 ];
2002 let hits = rrf_fuse(text, vec![], 10, "flashattention");
2003 assert_eq!(hits.len(), 2);
2004 assert_eq!(
2005 hits[0].entity_id, exact_id,
2006 "exact match must rank first despite being rank-2 in raw text search"
2007 );
2008 }
2009
2010 #[test]
2013 fn embed_batch_unconfigured_on_memory_runtime() {
2014 let rt = KhiveRuntime::memory().unwrap();
2016 let result = tokio::runtime::Runtime::new()
2017 .unwrap()
2018 .block_on(rt.embed_batch(&[]));
2019 assert!(result.is_ok());
2021 assert!(result.unwrap().is_empty());
2022 }
2023
2024 #[test]
2025 fn embed_batch_empty_input_returns_empty_vec() {
2026 let rt = KhiveRuntime::memory().unwrap();
2028 let result = tokio::runtime::Runtime::new()
2029 .unwrap()
2030 .block_on(rt.embed_batch(&[]));
2031 assert_eq!(result.unwrap(), Vec::<Vec<f32>>::new());
2032 }
2033
2034 #[test]
2035 fn embed_batch_no_model_non_empty_returns_unconfigured() {
2036 let rt = KhiveRuntime::memory().unwrap();
2037 let texts = vec!["hello".to_string()];
2038 let result = tokio::runtime::Runtime::new()
2039 .unwrap()
2040 .block_on(rt.embed_batch(&texts));
2041 match result {
2042 Err(crate::RuntimeError::Unconfigured(s)) => assert_eq!(s, "embedding_model"),
2043 Err(other) => panic!("expected Unconfigured, got {:?}", other),
2044 Ok(_) => panic!("expected Err, got Ok"),
2045 }
2046 }
2047
2048 #[test]
2049 #[ignore = "loads ~80 MB model; run with --include-ignored"]
2050 fn embed_batch_count_matches_input() {
2051 let config = RuntimeConfig {
2052 db_path: None,
2053 default_namespace: Namespace::parse("test").unwrap(),
2054 embedding_model: Some(EmbeddingModel::AllMiniLmL6V2),
2055 packs: vec!["kg".to_string()],
2056 ..RuntimeConfig::default()
2057 };
2058 let rt = KhiveRuntime::new(config).unwrap();
2059 let texts: Vec<String> = vec!["foo".to_string(), "bar".to_string(), "baz".to_string()];
2060 let result = tokio::runtime::Runtime::new()
2061 .unwrap()
2062 .block_on(rt.embed_batch(&texts));
2063 let embeddings = result.unwrap();
2064 assert_eq!(embeddings.len(), texts.len());
2065 }
2066
2067 #[test]
2068 fn vector_search_requires_embedding_or_text() {
2069 let rt = KhiveRuntime::memory().unwrap();
2070 let tok = NamespaceToken::local();
2071 let result = tokio::runtime::Runtime::new()
2072 .unwrap()
2073 .block_on(rt.vector_search(&tok, None, None, 10, Some(SubstrateKind::Entity)));
2074 match result {
2075 Err(crate::RuntimeError::InvalidInput(msg)) => {
2076 assert!(msg.contains("query_embedding or query_text"), "msg: {msg}");
2077 }
2078 other => panic!("expected InvalidInput, got {other:?}"),
2079 }
2080 }
2081
2082 #[test]
2083 fn vector_search_text_without_model_returns_unconfigured() {
2084 let rt = KhiveRuntime::memory().unwrap();
2085 let tok = NamespaceToken::local();
2086 let result = tokio::runtime::Runtime::new()
2087 .unwrap()
2088 .block_on(rt.vector_search(
2089 &tok,
2090 None,
2091 Some("attention"),
2092 10,
2093 Some(SubstrateKind::Entity),
2094 ));
2095 match result {
2096 Err(crate::RuntimeError::Unconfigured(s)) => assert_eq!(s, "embedding_model"),
2097 other => panic!("expected Unconfigured, got {other:?}"),
2098 }
2099 }
2100
2101 #[test]
2102 #[ignore = "loads ~80 MB model; run with --include-ignored"]
2103 fn embed_batch_vectors_have_expected_dimensions() {
2104 let model = EmbeddingModel::AllMiniLmL6V2;
2105 let config = RuntimeConfig {
2106 db_path: None,
2107 default_namespace: Namespace::parse("test").unwrap(),
2108 embedding_model: Some(model),
2109 packs: vec!["kg".to_string()],
2110 ..RuntimeConfig::default()
2111 };
2112 let rt = KhiveRuntime::new(config).unwrap();
2113 let texts = vec!["hello world".to_string()];
2114 let result = tokio::runtime::Runtime::new()
2115 .unwrap()
2116 .block_on(rt.embed_batch(&texts));
2117 let embeddings = result.unwrap();
2118 assert_eq!(embeddings[0].len(), model.dimensions());
2119 }
2120
2121 #[tokio::test]
2129 async fn hybrid_search_still_fails_loud_on_vector_arm_error() {
2130 let (rt, fail) = runtime_with_toggleable_embeddings();
2131 let tok = NamespaceToken::local();
2132 rt.create_entity(
2133 &tok,
2134 "concept",
2135 None,
2136 "FlashAttention",
2137 Some("IO-aware exact attention using tiling"),
2138 None,
2139 vec![],
2140 )
2141 .await
2142 .unwrap();
2143 break_vector_arm(&fail);
2144
2145 let result = rt
2146 .hybrid_search(&tok, "FlashAttention", None, 10, None, None, &[], None)
2147 .await;
2148
2149 assert!(
2150 result.is_err(),
2151 "the fail-loud entry point must still propagate a vector-arm failure, got {result:?}"
2152 );
2153 }
2154
2155 #[tokio::test]
2159 async fn hybrid_search_outcome_preserves_text_hits_on_vector_arm_error() {
2160 let (rt, fail) = runtime_with_toggleable_embeddings();
2161 let tok = NamespaceToken::local();
2162 rt.create_entity(
2163 &tok,
2164 "concept",
2165 None,
2166 "FlashAttention",
2167 Some("IO-aware exact attention using tiling"),
2168 None,
2169 vec![],
2170 )
2171 .await
2172 .unwrap();
2173 break_vector_arm(&fail);
2174
2175 let outcome = rt
2176 .hybrid_search_outcome(&tok, "FlashAttention", 10, None, None, &[], None)
2177 .await
2178 .expect("text leg must still succeed");
2179
2180 assert!(
2181 !outcome.hits.is_empty(),
2182 "text arm's hit must survive a vector-arm failure"
2183 );
2184 assert!(
2185 outcome.hits[0]
2186 .title
2187 .as_deref()
2188 .unwrap_or_default()
2189 .contains("FlashAttention"),
2190 "surviving hit must be the text match"
2191 );
2192 let vector_error = outcome
2193 .vector_error
2194 .expect("vector arm failure must be reported");
2195 assert!(
2196 vector_error.contains("injected vector-arm failure"),
2197 "vector_error must carry the underlying cause, got {vector_error:?}"
2198 );
2199 }
2200
2201 #[tokio::test]
2204 async fn hybrid_search_outcome_has_no_vector_error_when_vector_arm_healthy() {
2205 let rt = runtime_with_constant_embeddings();
2206 let tok = NamespaceToken::local();
2207 rt.create_entity(
2208 &tok,
2209 "concept",
2210 None,
2211 "FlashAttention",
2212 Some("IO-aware exact attention using tiling"),
2213 None,
2214 vec![],
2215 )
2216 .await
2217 .unwrap();
2218
2219 let outcome = rt
2220 .hybrid_search_outcome(&tok, "FlashAttention", 10, None, None, &[], None)
2221 .await
2222 .expect("hybrid search must succeed");
2223
2224 assert!(!outcome.hits.is_empty(), "should find the entity");
2225 assert!(
2226 outcome.vector_error.is_none(),
2227 "a healthy vector arm must not report an error"
2228 );
2229 }
2230
2231 #[tokio::test]
2234 async fn hybrid_search_entity_hit_has_title() {
2235 let rt = KhiveRuntime::memory().unwrap();
2236 let tok = NamespaceToken::local();
2237 rt.create_entity(
2238 &tok,
2239 "concept",
2240 None,
2241 "FlashAttention",
2242 Some("IO-aware exact attention using tiling"),
2243 None,
2244 vec![],
2245 )
2246 .await
2247 .unwrap();
2248
2249 let hits = rt
2250 .hybrid_search(&tok, "FlashAttention", None, 10, None, None, &[], None)
2251 .await
2252 .unwrap();
2253
2254 assert!(!hits.is_empty(), "should find the entity");
2255 let hit = &hits[0];
2256 assert!(hit.title.is_some(), "title must be populated");
2257 assert!(
2258 hit.title.as_deref().unwrap().contains("FlashAttention"),
2259 "title must contain entity name"
2260 );
2261 }
2262
2263 #[tokio::test]
2269 async fn hybrid_search_with_dollar_sign_query_does_not_error() {
2270 let rt = KhiveRuntime::memory().unwrap();
2271 let tok = NamespaceToken::local();
2272 rt.create_entity(
2273 &tok,
2274 "concept",
2275 None,
2276 "DSL docs",
2277 Some("use $prev.id to chain calls"),
2278 None,
2279 vec![],
2280 )
2281 .await
2282 .unwrap();
2283
2284 let result = rt
2285 .hybrid_search(&tok, "$prev.id", None, 10, None, None, &[], None)
2286 .await;
2287
2288 assert!(
2289 result.is_ok(),
2290 "#388 hybrid_search must not hard-fail on a '$'-bearing query, got: {:?}",
2291 result.err()
2292 );
2293 }
2294
2295 #[tokio::test]
2304 async fn hybrid_search_with_residual_fts5_char_now_sanitized() {
2305 let rt = KhiveRuntime::memory().unwrap();
2306 let tok = NamespaceToken::local();
2307 rt.create_entity(
2308 &tok,
2309 "concept",
2310 None,
2311 "DSL docs",
2312 Some("use foo@bar to chain calls"),
2313 None,
2314 vec![],
2315 )
2316 .await
2317 .unwrap();
2318
2319 let result = rt
2320 .hybrid_search(&tok, "foo@bar", None, 10, None, None, &[], None)
2321 .await;
2322
2323 let hits = result.unwrap_or_else(|e| {
2324 panic!("#916 hybrid_search must not fail on an '@'-bearing query, got: {e:?}")
2325 });
2326 assert!(
2327 !hits.is_empty(),
2328 "#916 '@'-bearing query must still find the seeded 'foo@bar' content via the \
2329 quoted-phrase alternative"
2330 );
2331 }
2332
2333 #[tokio::test]
2341 async fn hybrid_search_with_916_issue_characters_finds_text_leg_hits() {
2342 let rt = KhiveRuntime::memory().unwrap();
2343 let tok = NamespaceToken::local();
2344
2345 rt.create_entity(
2346 &tok,
2347 "concept",
2348 None,
2349 "issue tracker",
2350 Some("tracking #682 Stage 2: MoE expert-cache prefetch work"),
2351 None,
2352 vec![],
2353 )
2354 .await
2355 .unwrap();
2356 rt.create_entity(
2357 &tok,
2358 "concept",
2359 None,
2360 "benchmark notes",
2361 Some("chunkwise B=128 traffic arithmetic simdgroup_matrix DPLR"),
2362 None,
2363 vec![],
2364 )
2365 .await
2366 .unwrap();
2367 rt.create_entity(
2368 &tok,
2369 "concept",
2370 None,
2371 "sampling notes",
2372 Some("evaluated with the Min-K%Prob membership inference method"),
2373 None,
2374 vec![],
2375 )
2376 .await
2377 .unwrap();
2378
2379 for query in ["#682 Stage 2", "B=128", "Min-K%Prob"] {
2380 let result = rt
2381 .hybrid_search(&tok, query, None, 10, None, None, &[], None)
2382 .await;
2383 let hits = result.unwrap_or_else(|e| {
2384 panic!("#916 hybrid_search must not fail on query {query:?}, got: {e:?}")
2385 });
2386 assert!(
2387 hits.iter()
2388 .any(|h| matches!(h.source, SearchSource::Text | SearchSource::Both)),
2389 "#916 query {query:?} must surface a Text/Both-sourced hit \
2390 (the FTS leg must contribute, not just the vector leg); got {hits:?}"
2391 );
2392 }
2393 }
2394
2395 #[tokio::test]
2409 async fn hybrid_search_tag_filter_pushed_before_truncation() {
2410 let rt = KhiveRuntime::memory().unwrap();
2411 let tok = NamespaceToken::local();
2412
2413 rt.create_entity(
2415 &tok,
2416 "concept",
2417 None,
2418 "alpha beta gamma decoy alpha beta gamma",
2419 Some("alpha beta gamma decoy description alpha beta gamma"),
2420 None,
2421 vec!["other-tag".to_string()],
2422 )
2423 .await
2424 .unwrap();
2425
2426 let target = rt
2428 .create_entity(
2429 &tok,
2430 "concept",
2431 None,
2432 "alpha beta gamma target",
2433 Some("alpha beta gamma target description"),
2434 None,
2435 vec!["target-tag".to_string()],
2436 )
2437 .await
2438 .unwrap();
2439
2440 let hits = rt
2444 .hybrid_search(
2445 &tok,
2446 "alpha beta gamma",
2447 None,
2448 1,
2449 None,
2450 None,
2451 &["target-tag".to_string()],
2452 None,
2453 )
2454 .await
2455 .unwrap();
2456
2457 assert_eq!(
2458 hits.len(),
2459 1,
2460 "exactly one hit expected (the tag-matching entity)"
2461 );
2462 assert_eq!(
2463 hits[0].entity_id, target.id,
2464 "the tag-filtered entity must be returned even when ranked below limit in raw fusion"
2465 );
2466 }
2467
2468 #[tokio::test]
2477 async fn hybrid_search_props_filter_pushed_before_truncation() {
2478 let rt = KhiveRuntime::memory().unwrap();
2479 let tok = NamespaceToken::local();
2480
2481 rt.create_entity(
2482 &tok,
2483 "concept",
2484 None,
2485 "delta epsilon zeta decoy delta epsilon zeta",
2486 Some("delta epsilon zeta decoy description delta epsilon zeta"),
2487 Some(serde_json::json!({"domain": "other"})),
2488 vec![],
2489 )
2490 .await
2491 .unwrap();
2492
2493 let target = rt
2494 .create_entity(
2495 &tok,
2496 "concept",
2497 None,
2498 "delta epsilon zeta target",
2499 Some("delta epsilon zeta target description"),
2500 Some(serde_json::json!({"domain": "target"})),
2501 vec![],
2502 )
2503 .await
2504 .unwrap();
2505
2506 let filter = serde_json::json!({"domain": "target"});
2507 let hits = rt
2508 .hybrid_search(
2509 &tok,
2510 "delta epsilon zeta",
2511 None,
2512 1,
2513 None,
2514 None,
2515 &[],
2516 Some(&filter),
2517 )
2518 .await
2519 .unwrap();
2520
2521 assert_eq!(hits.len(), 1, "exactly one hit expected (properties match)");
2522 assert_eq!(
2523 hits[0].entity_id, target.id,
2524 "the properties-filtered entity must be returned even when ranked below limit"
2525 );
2526 }
2527
2528 #[tokio::test]
2538 async fn hybrid_search_entity_kind_filter_pushed_into_text_arm() {
2539 let rt = KhiveRuntime::memory().unwrap();
2540 let tok = NamespaceToken::local();
2541
2542 for i in 0..12 {
2543 rt.create_entity(
2544 &tok,
2545 "concept",
2546 None,
2547 &format!("quillfeather decoy {i}"),
2548 Some("quillfeather quillfeather quillfeather"),
2549 None,
2550 vec![],
2551 )
2552 .await
2553 .unwrap();
2554 }
2555
2556 let description = format!(
2557 "A long administrative description that mentions quillfeather once. {}",
2558 "Unrelated filing, scheduling and review words. ".repeat(20)
2559 );
2560 let target = rt
2561 .create_entity(
2562 &tok,
2563 "document",
2564 None,
2565 "Archive filing report",
2566 Some(description.as_str()),
2567 None,
2568 vec![],
2569 )
2570 .await
2571 .unwrap();
2572
2573 let unfiltered = rt
2576 .hybrid_search(&tok, "quillfeather", None, 13, None, None, &[], None)
2577 .await
2578 .unwrap();
2579 let target_rank = unfiltered
2580 .iter()
2581 .position(|hit| hit.entity_id == target.id)
2582 .expect("the document is found without a kind filter");
2583 assert!(
2584 target_rank >= 4,
2585 "the document must rank behind the candidate window, got rank {target_rank}"
2586 );
2587
2588 let hits = rt
2589 .hybrid_search(
2590 &tok,
2591 "quillfeather",
2592 None,
2593 1,
2594 Some("document"),
2595 None,
2596 &[],
2597 None,
2598 )
2599 .await
2600 .unwrap();
2601
2602 assert_eq!(
2603 hits.len(),
2604 1,
2605 "the kind-filtered search must return the matching document"
2606 );
2607 assert_eq!(
2608 hits[0].entity_id, target.id,
2609 "the returned hit must be the document, not a concept"
2610 );
2611 }
2612
2613 #[tokio::test]
2621 async fn hybrid_search_each_kind_matches_per_kind_search_when_one_kind_dominates() {
2622 let rt = KhiveRuntime::memory().unwrap();
2623 let tok = NamespaceToken::local();
2624
2625 for i in 0..20 {
2626 rt.create_entity(
2627 &tok,
2628 "concept",
2629 None,
2630 &format!("quillfeather decoy {i}"),
2631 Some("quillfeather quillfeather quillfeather"),
2632 None,
2633 vec![],
2634 )
2635 .await
2636 .unwrap();
2637 }
2638 for i in 0..3 {
2639 let description = format!(
2640 "A long administrative description {i} that mentions quillfeather once. {}",
2641 "Unrelated filing, scheduling and review words. ".repeat(20)
2642 );
2643 rt.create_entity(
2644 &tok,
2645 "document",
2646 None,
2647 &format!("Archive filing report {i}"),
2648 Some(description.as_str()),
2649 None,
2650 vec![],
2651 )
2652 .await
2653 .unwrap();
2654 }
2655
2656 let kinds = ["concept", "document"];
2657 let mut expected: Vec<Vec<Uuid>> = Vec::new();
2658 for kind in kinds {
2659 let hits = rt
2660 .hybrid_search(&tok, "quillfeather", None, 2, Some(kind), None, &[], None)
2661 .await
2662 .unwrap();
2663 expected.push(hits.iter().map(|hit| hit.entity_id).collect());
2664 }
2665 assert_eq!(
2666 expected[0].len(),
2667 2,
2668 "premise: the concept list is cut to limit"
2669 );
2670 assert_eq!(
2671 expected[1].len(),
2672 2,
2673 "premise: the document list is not starved"
2674 );
2675
2676 let per_kind = rt
2677 .hybrid_search_each_kind(&tok, "quillfeather", None, 2, &kinds)
2678 .await
2679 .unwrap();
2680 let mut actual: Vec<Vec<Uuid>> = Vec::new();
2681 for hits in &per_kind {
2682 actual.push(hits.iter().map(|hit| hit.entity_id).collect());
2683 }
2684 assert_eq!(
2685 actual, expected,
2686 "each kind must get the ids and order its own per-kind search returns"
2687 );
2688 }
2689
2690 struct CapturingEmbeddingService {
2693 captured: std::sync::Arc<std::sync::Mutex<Vec<Vec<String>>>>,
2694 }
2695
2696 #[async_trait::async_trait]
2697 impl EmbeddingService for CapturingEmbeddingService {
2698 async fn embed(
2699 &self,
2700 texts: &[String],
2701 _model: EmbeddingModel,
2702 ) -> std::result::Result<Vec<Vec<f32>>, lattice_embed::EmbedError> {
2703 self.captured.lock().unwrap().push(texts.to_vec());
2704 Ok(texts.iter().map(|_| vec![1.0]).collect())
2705 }
2706
2707 fn supports_model(&self, _model: EmbeddingModel) -> bool {
2708 true
2709 }
2710
2711 fn name(&self) -> &'static str {
2712 "capturing-embedding-service"
2713 }
2714 }
2715
2716 struct CapturingEmbedderProvider {
2717 name: String,
2718 captured: std::sync::Arc<std::sync::Mutex<Vec<Vec<String>>>>,
2719 }
2720
2721 struct RewritingPassageService {
2722 captured: std::sync::Arc<std::sync::Mutex<Vec<Vec<String>>>>,
2723 }
2724
2725 #[async_trait::async_trait]
2726 impl EmbeddingService for RewritingPassageService {
2727 async fn embed(
2728 &self,
2729 texts: &[String],
2730 _model: EmbeddingModel,
2731 ) -> std::result::Result<Vec<Vec<f32>>, lattice_embed::EmbedError> {
2732 self.captured.lock().unwrap().push(texts.to_vec());
2733 Ok(texts.iter().map(|_| vec![1.0]).collect())
2734 }
2735
2736 async fn embed_passage(
2737 &self,
2738 texts: &[String],
2739 model: EmbeddingModel,
2740 ) -> std::result::Result<Vec<Vec<f32>>, lattice_embed::EmbedError> {
2741 let prepared: Vec<String> = texts.iter().map(|text| format!("custom:{text}")).collect();
2742 self.embed(&prepared, model).await
2743 }
2744
2745 fn supports_model(&self, _model: EmbeddingModel) -> bool {
2746 true
2747 }
2748
2749 fn name(&self) -> &'static str {
2750 "rewriting-passage-service"
2751 }
2752 }
2753
2754 struct RewritingPassageProvider {
2755 name: String,
2756 captured: std::sync::Arc<std::sync::Mutex<Vec<Vec<String>>>>,
2757 }
2758
2759 #[async_trait::async_trait]
2760 impl EmbedderProvider for RewritingPassageProvider {
2761 fn name(&self) -> &str {
2762 &self.name
2763 }
2764
2765 fn dimensions(&self) -> usize {
2766 1
2767 }
2768
2769 async fn build(&self) -> crate::error::RuntimeResult<std::sync::Arc<dyn EmbeddingService>> {
2770 Ok(std::sync::Arc::new(RewritingPassageService {
2771 captured: std::sync::Arc::clone(&self.captured),
2772 }))
2773 }
2774 }
2775
2776 struct WrongCardinalityEmbeddingService;
2777
2778 #[async_trait::async_trait]
2779 impl EmbeddingService for WrongCardinalityEmbeddingService {
2780 async fn embed(
2781 &self,
2782 texts: &[String],
2783 _model: EmbeddingModel,
2784 ) -> std::result::Result<Vec<Vec<f32>>, lattice_embed::EmbedError> {
2785 Ok(texts
2786 .iter()
2787 .take(texts.len().saturating_sub(1))
2788 .map(|_| vec![1.0])
2789 .collect())
2790 }
2791
2792 fn supports_model(&self, _model: EmbeddingModel) -> bool {
2793 true
2794 }
2795
2796 fn name(&self) -> &'static str {
2797 "wrong-cardinality-embedding-service"
2798 }
2799 }
2800
2801 struct WrongCardinalityEmbedderProvider;
2802
2803 #[async_trait::async_trait]
2804 impl EmbedderProvider for WrongCardinalityEmbedderProvider {
2805 fn name(&self) -> &str {
2806 "wrong-cardinality-embedding-service"
2807 }
2808
2809 fn dimensions(&self) -> usize {
2810 1
2811 }
2812
2813 async fn build(&self) -> crate::error::RuntimeResult<std::sync::Arc<dyn EmbeddingService>> {
2814 Ok(std::sync::Arc::new(WrongCardinalityEmbeddingService))
2815 }
2816 }
2817
2818 struct SurplusCardinalityEmbeddingService;
2819
2820 #[async_trait::async_trait]
2821 impl EmbeddingService for SurplusCardinalityEmbeddingService {
2822 async fn embed(
2823 &self,
2824 texts: &[String],
2825 _model: EmbeddingModel,
2826 ) -> std::result::Result<Vec<Vec<f32>>, lattice_embed::EmbedError> {
2827 let mut vectors: Vec<Vec<f32>> = texts.iter().map(|_| vec![1.0]).collect();
2828 vectors.push(vec![1.0]);
2829 Ok(vectors)
2830 }
2831
2832 fn supports_model(&self, _model: EmbeddingModel) -> bool {
2833 true
2834 }
2835
2836 fn name(&self) -> &'static str {
2837 "surplus-cardinality-embedding-service"
2838 }
2839 }
2840
2841 struct SurplusCardinalityEmbedderProvider;
2842
2843 #[async_trait::async_trait]
2844 impl EmbedderProvider for SurplusCardinalityEmbedderProvider {
2845 fn name(&self) -> &str {
2846 "surplus-cardinality-embedding-service"
2847 }
2848
2849 fn dimensions(&self) -> usize {
2850 1
2851 }
2852
2853 async fn build(&self) -> crate::error::RuntimeResult<std::sync::Arc<dyn EmbeddingService>> {
2854 Ok(std::sync::Arc::new(SurplusCardinalityEmbeddingService))
2855 }
2856 }
2857
2858 #[async_trait::async_trait]
2859 impl EmbedderProvider for CapturingEmbedderProvider {
2860 fn name(&self) -> &str {
2861 &self.name
2862 }
2863
2864 fn dimensions(&self) -> usize {
2865 1
2866 }
2867
2868 async fn build(&self) -> crate::error::RuntimeResult<std::sync::Arc<dyn EmbeddingService>> {
2869 Ok(std::sync::Arc::new(CapturingEmbeddingService {
2870 captured: std::sync::Arc::clone(&self.captured),
2871 }))
2872 }
2873 }
2874
2875 fn runtime_with_capturing_embedder(
2876 model: EmbeddingModel,
2877 ) -> (
2878 KhiveRuntime,
2879 std::sync::Arc<std::sync::Mutex<Vec<Vec<String>>>>,
2880 ) {
2881 let runtime = KhiveRuntime::memory().unwrap();
2882 let captured = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
2883 runtime.register_embedder(CapturingEmbedderProvider {
2884 name: model.to_string(),
2885 captured: std::sync::Arc::clone(&captured),
2886 });
2887 (runtime, captured)
2888 }
2889
2890 #[tokio::test]
2891 async fn bge_query_paths_pass_raw_unprefixed_text() {
2892 const BGE_QUERY_INSTRUCTION: &str =
2893 "Represent this sentence for searching relevant passages: ";
2894 let single = "single raw query";
2895 let batch = vec![
2896 "first raw query".to_string(),
2897 "second raw query".to_string(),
2898 ];
2899
2900 for model in [
2901 EmbeddingModel::BgeSmallEnV15,
2902 EmbeddingModel::BgeBaseEnV15,
2903 EmbeddingModel::BgeLargeEnV15,
2904 ] {
2905 let (runtime, captured) = runtime_with_capturing_embedder(model);
2906 runtime
2907 .embed_query_with_model(&model.to_string(), single)
2908 .await
2909 .unwrap();
2910 runtime
2911 .embed_query_batch_with_model(&model.to_string(), &batch)
2912 .await
2913 .unwrap();
2914
2915 let calls = captured.lock().unwrap().clone();
2916 assert_eq!(
2917 calls,
2918 vec![vec![single.to_string()], batch.clone()],
2919 "{model} must receive raw query text through single and batch paths"
2920 );
2921 assert!(
2922 calls
2923 .iter()
2924 .flatten()
2925 .all(|text| !text.contains(BGE_QUERY_INSTRUCTION)),
2926 "{model} must not receive the BGE retrieval instruction"
2927 );
2928 }
2929 }
2930
2931 #[tokio::test]
2932 async fn e5_query_paths_apply_query_prefix() {
2933 let model = EmbeddingModel::MultilingualE5Small;
2934 let single = "single raw query";
2935 let batch = vec![
2936 "first raw query".to_string(),
2937 "second raw query".to_string(),
2938 ];
2939 let (runtime, captured) = runtime_with_capturing_embedder(model);
2940
2941 runtime
2942 .embed_query_with_model(&model.to_string(), single)
2943 .await
2944 .unwrap();
2945 runtime
2946 .embed_query_batch_with_model(&model.to_string(), &batch)
2947 .await
2948 .unwrap();
2949
2950 assert_eq!(
2951 captured.lock().unwrap().as_slice(),
2952 [
2953 vec!["query: single raw query".to_string()],
2954 vec![
2955 "query: first raw query".to_string(),
2956 "query: second raw query".to_string(),
2957 ],
2958 ],
2959 "E5 must receive its query prefix through single and batch paths"
2960 );
2961 }
2962
2963 #[tokio::test]
2964 async fn mixed_document_batch_stays_one_ordered_provider_call() {
2965 let model = EmbeddingModel::AllMiniLmL6V2;
2966 let (runtime, captured) = runtime_with_capturing_embedder(model);
2967 let texts = vec![
2968 "first normal document".to_string(),
2969 "x".repeat(MAX_TEXT_BYTES + 1),
2970 "second normal document".to_string(),
2971 ];
2972
2973 let outcomes = runtime
2974 .embed_document_batch_with_model_outcomes(&model.to_string(), &texts)
2975 .await
2976 .expect("mixed batch must embed");
2977 assert_eq!(outcomes.len(), texts.len());
2978 assert!(!outcomes[0].truncated);
2979 assert_eq!(outcomes[0].source_bytes, outcomes[0].embedded_bytes);
2980 assert!(outcomes[1].truncated);
2981 assert_eq!(outcomes[1].source_bytes, MAX_TEXT_BYTES + 1);
2982 assert_eq!(outcomes[1].embedded_bytes, MAX_TEXT_BYTES);
2983 assert!(!outcomes[2].truncated);
2984 assert_eq!(
2985 captured.lock().unwrap().as_slice(),
2986 [vec![
2987 texts[0].clone(),
2988 "x".repeat(MAX_TEXT_BYTES),
2989 texts[2].clone(),
2990 ]]
2991 );
2992 }
2993
2994 #[tokio::test]
2995 async fn canonical_name_override_has_no_exact_input_fingerprint() {
2996 let model = EmbeddingModel::MultilingualE5Small;
2997 let model_name = model.to_string();
2998 let (runtime, captured) = runtime_with_capturing_embedder(model);
2999 let input = "x".repeat(document_embedding_budget(&model_name) + 3);
3000 let outcomes = runtime
3001 .embed_document_batch_with_model_outcomes(&model_name, std::slice::from_ref(&input))
3002 .await
3003 .unwrap();
3004 let calls = captured.lock().unwrap();
3005 let actual_input = &calls[0][0];
3006 assert!(actual_input.starts_with("passage: "));
3007 assert_eq!(actual_input.len(), MAX_TEXT_BYTES);
3008 assert!(outcomes[0].truncated);
3009 assert_eq!(outcomes[0].prepared_text_fingerprint, None);
3010 assert_eq!(
3011 prepared_document_fingerprint(
3012 bounded_embedding_input(&input, document_embedding_budget(&model_name)).0,
3013 model,
3014 ),
3015 VectorRecord::fingerprint_text(actual_input),
3016 "the audited lattice path hashes the bounded text plus passage prefix"
3017 );
3018 }
3019
3020 #[tokio::test]
3021 async fn ordinary_custom_document_provider_has_no_exact_input_fingerprint() {
3022 let runtime = KhiveRuntime::memory().unwrap();
3023 let captured = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
3024 runtime.register_embedder(CapturingEmbedderProvider {
3025 name: "custom-document-provider".to_string(),
3026 captured: std::sync::Arc::clone(&captured),
3027 });
3028
3029 let outcome = runtime
3030 .embed_document_with_model_outcome("custom-document-provider", "document")
3031 .await
3032 .expect("custom provider must embed");
3033 assert_eq!(outcome.prepared_text_fingerprint, None);
3034 assert_eq!(captured.lock().unwrap().len(), 1);
3035 }
3036
3037 #[tokio::test]
3038 async fn custom_passage_override_under_builtin_name_cannot_claim_exact_input() {
3039 let runtime = KhiveRuntime::memory().unwrap();
3040 let captured = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
3041 let model = EmbeddingModel::MultilingualE5Small;
3042 let model_name = model.to_string();
3043 runtime.register_embedder(RewritingPassageProvider {
3044 name: model_name.clone(),
3045 captured: std::sync::Arc::clone(&captured),
3046 });
3047
3048 let outcome = runtime
3049 .embed_document_with_model_outcome(&model_name, "document")
3050 .await
3051 .expect("custom override must embed");
3052 assert_eq!(outcome.prepared_text_fingerprint, None);
3053 assert_eq!(
3054 captured.lock().unwrap()[0],
3055 vec!["custom:document".to_string()]
3056 );
3057 assert_ne!(
3058 prepared_document_fingerprint("document", model),
3059 VectorRecord::fingerprint_text("custom:document")
3060 );
3061 }
3062
3063 #[tokio::test]
3064 async fn truncated_document_batch_rejects_provider_cardinality_mismatch() {
3065 let runtime = KhiveRuntime::memory().unwrap();
3066 runtime.register_embedder(WrongCardinalityEmbedderProvider);
3067 let model_name = "wrong-cardinality-embedding-service";
3068 let texts = vec!["x".repeat(MAX_TEXT_BYTES + 1), "normal".to_string()];
3069
3070 let error = runtime
3071 .embed_document_batch_with_model_outcomes(model_name, &texts)
3072 .await
3073 .expect_err("provider cardinality mismatch must fail the whole batch");
3074
3075 assert!(
3076 error
3077 .to_string()
3078 .contains("embed_passage returned 1 vectors for 2 inputs"),
3079 "unexpected error: {error}"
3080 );
3081 }
3082
3083 #[tokio::test]
3084 async fn singleton_document_embed_rejects_zero_vectors() {
3085 let runtime = KhiveRuntime::memory().unwrap();
3086 runtime.register_embedder(WrongCardinalityEmbedderProvider);
3087 let model_name = "wrong-cardinality-embedding-service";
3088
3089 let error = runtime
3090 .embed_document_with_model_outcome(model_name, "single document")
3091 .await
3092 .expect_err("provider returning zero vectors must fail closed");
3093
3094 assert!(
3095 error
3096 .to_string()
3097 .contains("embed_passage returned 0 vectors for 1 input"),
3098 "unexpected error: {error}"
3099 );
3100 }
3101
3102 #[tokio::test]
3103 async fn singleton_document_embed_rejects_surplus_vectors() {
3104 let runtime = KhiveRuntime::memory().unwrap();
3105 runtime.register_embedder(SurplusCardinalityEmbedderProvider);
3106 let model_name = "surplus-cardinality-embedding-service";
3107
3108 let error = runtime
3109 .embed_document_with_model_outcome(model_name, "single document")
3110 .await
3111 .expect_err("provider returning surplus vectors must fail closed");
3112
3113 assert!(
3114 error
3115 .to_string()
3116 .contains("embed_passage returned 2 vectors for 1 input"),
3117 "unexpected error: {error}"
3118 );
3119 }
3120
3121 #[test]
3122 #[ignore = "loads ~80 MB model; run with --include-ignored"]
3123 fn minilm_document_and_query_embed_are_identical_no_prefix_model() {
3124 let model = EmbeddingModel::AllMiniLmL6V2;
3128 let config = RuntimeConfig {
3129 db_path: None,
3130 default_namespace: Namespace::parse("test").unwrap(),
3131 embedding_model: Some(model),
3132 packs: vec!["kg".to_string()],
3133 ..RuntimeConfig::default()
3134 };
3135 let rt = KhiveRuntime::new(config).unwrap();
3136 let text = "attention is all you need".to_string();
3137 let rt_ref = &rt;
3138 let (doc_emb, query_emb) = tokio::runtime::Runtime::new().unwrap().block_on(async {
3139 let d = rt_ref
3140 .embed_document_with_model_outcome(&model.to_string(), &text)
3141 .await
3142 .unwrap()
3143 .vector;
3144 let q = rt_ref
3145 .embed_query_with_model(&model.to_string(), &text)
3146 .await
3147 .unwrap();
3148 (d, q)
3149 });
3150 assert_eq!(
3151 doc_emb, query_emb,
3152 "MiniLM has no instruction prefix: document and query embeds must be identical"
3153 );
3154 }
3155
3156 #[test]
3157 #[ignore = "loads multilingual-e5-small (~90 MB); run with --include-ignored"]
3158 fn e5_document_and_query_embed_differ_instruction_tuned_model() {
3159 let model = EmbeddingModel::MultilingualE5Small;
3164 let config = RuntimeConfig {
3165 db_path: None,
3166 default_namespace: Namespace::parse("test").unwrap(),
3167 embedding_model: Some(model),
3168 packs: vec!["kg".to_string()],
3169 ..RuntimeConfig::default()
3170 };
3171 let rt = KhiveRuntime::new(config).unwrap();
3172 let text = "attention is all you need".to_string();
3173 let rt_ref = &rt;
3174 let (doc_emb, query_emb) = tokio::runtime::Runtime::new().unwrap().block_on(async {
3175 let d = rt_ref
3176 .embed_document_with_model_outcome(&model.to_string(), &text)
3177 .await
3178 .unwrap()
3179 .vector;
3180 let q = rt_ref
3181 .embed_query_with_model(&model.to_string(), &text)
3182 .await
3183 .unwrap();
3184 (d, q)
3185 });
3186 assert_ne!(
3187 doc_emb, query_emb,
3188 "multilingual-e5-small uses asymmetric prefixes: document ('passage: ') \
3189 and query ('query: ') embeds of the same text must differ"
3190 );
3191 }
3192
3193 use crate::embedder_registry::EmbedderProvider;
3196 use lattice_embed::EmbeddingService;
3197
3198 struct StubEmbedderProvider;
3203
3204 #[async_trait::async_trait]
3205 impl EmbedderProvider for StubEmbedderProvider {
3206 fn name(&self) -> &str {
3207 "stub-model-m07"
3208 }
3209
3210 fn dimensions(&self) -> usize {
3211 4
3212 }
3213
3214 async fn build(&self) -> crate::error::RuntimeResult<std::sync::Arc<dyn EmbeddingService>> {
3215 struct StubSvc;
3216 #[async_trait::async_trait]
3217 impl EmbeddingService for StubSvc {
3218 async fn embed(
3219 &self,
3220 _texts: &[String],
3221 _model: lattice_embed::EmbeddingModel,
3222 ) -> std::result::Result<Vec<Vec<f32>>, lattice_embed::EmbedError> {
3223 Ok(vec![])
3224 }
3225
3226 fn supports_model(&self, _model: lattice_embed::EmbeddingModel) -> bool {
3227 true
3228 }
3229
3230 fn name(&self) -> &'static str {
3231 "stub-svc-m07"
3232 }
3233 }
3234 Ok(std::sync::Arc::new(StubSvc))
3235 }
3236 }
3237
3238 #[tokio::test]
3246 async fn backfill_reader_error_is_propagated_not_swallowed() {
3247 let rt = KhiveRuntime::memory().unwrap();
3248 rt.register_embedder(StubEmbedderProvider);
3249 let tok = NamespaceToken::local();
3250
3251 super::arm_backfill_reader_fail();
3254
3255 let result = rt.backfill_missing_embeddings(&tok).await;
3256 assert!(
3257 result.is_err(),
3258 "backfill_missing_embeddings must propagate the reader error (got Ok instead)"
3259 );
3260 let err_msg = result.unwrap_err().to_string();
3261 assert!(
3262 err_msg.contains("injected failure"),
3263 "error must originate from the injected reader failure, got: {err_msg}"
3264 );
3265 }
3266
3267 #[tokio::test]
3268 async fn backfill_skips_a_full_page_of_excluded_messages_and_reaches_eligible_tail() {
3269 use crate::{NoteEmbeddingPolicy, NoteEmbeddingPolicySpec};
3270 use khive_storage::note::Note;
3271
3272 const PAGE: u128 = EMBEDDING_BATCH_PAGE_SIZE as u128;
3273 let primary = EmbeddingModel::AllMiniLmL6V2;
3274 let primary_name = primary.to_string();
3275 let secondary_name = "zz-backfill-secondary";
3276 let rt = KhiveRuntime::new(RuntimeConfig {
3277 db_path: None,
3278 embedding_model: Some(primary),
3279 packs: vec![],
3280 ..RuntimeConfig::no_embeddings()
3281 })
3282 .unwrap();
3283 rt.register_embedder(ConstantEmbedderProvider {
3284 name: primary_name.clone(),
3285 dimensions: primary.dimensions(),
3286 fail: Arc::new(AtomicBool::new(false)),
3287 });
3288 rt.register_embedder(ConstantEmbedderProvider {
3289 name: secondary_name.into(),
3290 dimensions: 4,
3291 fail: Arc::new(AtomicBool::new(false)),
3292 });
3293 rt.install_note_embedding_policies(&[NoteEmbeddingPolicySpec {
3294 kind: "message",
3295 policy: NoteEmbeddingPolicy::DefaultModel,
3296 }]);
3297 let tok = NamespaceToken::local();
3298 let primary_store = rt.vectors_for_model(&tok, &primary_name).unwrap();
3299 let secondary_store = rt.vectors_for_model(&tok, secondary_name).unwrap();
3300 let notes = rt.notes(&tok).unwrap();
3301 let mut seeded = Vec::new();
3302 for ordinal in 1..=PAGE + 1 {
3303 let mut message = Note::new("local", "message", "excluded from secondary");
3304 message.id = Uuid::from_u128(ordinal);
3305 seeded.push(message);
3306 }
3307 let mut ordinary = Note::new("local", "observation", "eligible after full page");
3308 ordinary.id = Uuid::from_u128(u128::MAX);
3309 seeded.push(ordinary.clone());
3310 notes.upsert_notes(seeded).await.unwrap();
3311
3312 let backfilled = tokio::time::timeout(
3313 std::time::Duration::from_secs(60),
3314 rt.backfill_missing_embeddings(&tok),
3315 )
3316 .await
3317 .expect("backfill must advance past a full excluded page")
3318 .unwrap();
3319 assert_eq!(backfilled, (PAGE + 2) as u64 + 1);
3320 assert_eq!(primary_store.count().await.unwrap(), (PAGE + 2) as u64);
3321 assert_eq!(secondary_store.count().await.unwrap(), 1);
3322 let excluded_document = rt
3323 .text_for_notes(&tok)
3324 .unwrap()
3325 .get_document("local", Uuid::from_u128(1))
3326 .await
3327 .unwrap()
3328 .expect("first-model pass must index an excluded message");
3329 assert_eq!(excluded_document.body, "excluded from secondary");
3330 assert!(
3331 rt.text_for_notes(&tok)
3332 .unwrap()
3333 .get_document("local", ordinary.id)
3334 .await
3335 .unwrap()
3336 .is_some(),
3337 "the first-model pass must still repopulate FTS"
3338 );
3339 }
3340
3341 const BACKFILL_BATCH_MODEL: &str = "backfill-batch-model";
3342
3343 struct CountingBackfillService {
3344 calls: Arc<AtomicUsize>,
3345 reject_poison: bool,
3346 }
3347
3348 #[async_trait::async_trait]
3349 impl EmbeddingService for CountingBackfillService {
3350 async fn embed(
3351 &self,
3352 texts: &[String],
3353 _model: EmbeddingModel,
3354 ) -> std::result::Result<Vec<Vec<f32>>, EmbedError> {
3355 self.calls.fetch_add(1, Ordering::SeqCst);
3356 if self.reject_poison
3357 && (texts.len() > 1 || texts.iter().any(|text| text.contains("poison")))
3358 {
3359 return Err(EmbedError::InferenceFailed("poison input".into()));
3360 }
3361 Ok(texts
3362 .iter()
3363 .map(|text| {
3364 vec![
3365 text.len() as f32,
3366 text.bytes().map(u32::from).sum::<u32>() as f32,
3367 text.as_bytes().first().copied().unwrap_or_default() as f32,
3368 text.as_bytes().last().copied().unwrap_or_default() as f32,
3369 ]
3370 })
3371 .collect())
3372 }
3373
3374 fn supports_model(&self, _model: EmbeddingModel) -> bool {
3375 true
3376 }
3377
3378 fn name(&self) -> &'static str {
3379 "backfill-batch-counting-service"
3380 }
3381 }
3382
3383 struct CountingBackfillProvider {
3384 calls: Arc<AtomicUsize>,
3385 reject_poison: bool,
3386 }
3387
3388 #[async_trait::async_trait]
3389 impl EmbedderProvider for CountingBackfillProvider {
3390 fn name(&self) -> &str {
3391 BACKFILL_BATCH_MODEL
3392 }
3393
3394 fn dimensions(&self) -> usize {
3395 4
3396 }
3397
3398 async fn build(&self) -> crate::error::RuntimeResult<Arc<dyn EmbeddingService>> {
3399 Ok(Arc::new(CountingBackfillService {
3400 calls: Arc::clone(&self.calls),
3401 reject_poison: self.reject_poison,
3402 }))
3403 }
3404 }
3405
3406 #[tokio::test]
3407 async fn backfill_batches_provider_calls_and_matches_single_record_vectors() {
3408 let token = NamespaceToken::local();
3409 let runtime = KhiveRuntime::memory().unwrap();
3410 let mut entity_ids = Vec::new();
3411 for index in 0..257 {
3412 let entity = runtime
3413 .create_entity(
3414 &token,
3415 "concept",
3416 None,
3417 &format!("Backfill entity {index}"),
3418 (index != 0).then_some("body"),
3419 None,
3420 vec![],
3421 )
3422 .await
3423 .unwrap();
3424 entity_ids.push(entity.id);
3425 }
3426 let mut note_ids = Vec::new();
3427 for index in 0..3 {
3428 let note = runtime
3429 .create_note(
3430 &token,
3431 "observation",
3432 None,
3433 &format!("Backfill note {index}"),
3434 None,
3435 None,
3436 vec![],
3437 )
3438 .await
3439 .unwrap();
3440 note_ids.push(note.id);
3441 }
3442 let calls = Arc::new(AtomicUsize::new(0));
3443 runtime.register_embedder(CountingBackfillProvider {
3444 calls: Arc::clone(&calls),
3445 reject_poison: false,
3446 });
3447 let usage = crate::usage::UsageContext::new();
3448 let backfilled =
3449 crate::usage::scope(usage.clone(), runtime.backfill_missing_embeddings(&token))
3450 .await
3451 .unwrap();
3452 assert_eq!(backfilled, 260);
3453 assert_eq!(
3454 calls.load(Ordering::SeqCst),
3455 3,
3456 "two entity pages and one note page"
3457 );
3458 assert_eq!(
3459 usage.snapshot()["embed_calls"],
3460 260,
3461 "ADR-103 counts texts, not provider calls"
3462 );
3463
3464 let singles = KhiveRuntime::memory().unwrap();
3465 let single_calls = Arc::new(AtomicUsize::new(0));
3466 singles.register_embedder(CountingBackfillProvider {
3467 calls: Arc::clone(&single_calls),
3468 reject_poison: false,
3469 });
3470 for id in &entity_ids {
3471 let entity = runtime.get_entity(&token, *id).await.unwrap();
3472 singles
3473 .entities(&token)
3474 .unwrap()
3475 .upsert_entity(entity.clone())
3476 .await
3477 .unwrap();
3478 singles.reindex_entity(&token, &entity).await.unwrap();
3479 }
3480 for id in ¬e_ids {
3481 let note = runtime
3482 .notes(&token)
3483 .unwrap()
3484 .get_note(*id)
3485 .await
3486 .unwrap()
3487 .unwrap();
3488 singles
3489 .notes(&token)
3490 .unwrap()
3491 .upsert_note(note.clone())
3492 .await
3493 .unwrap();
3494 singles.reindex_note(&token, ¬e).await.unwrap();
3495 }
3496 assert_eq!(single_calls.load(Ordering::SeqCst), 260);
3497 let batched_entities = runtime
3498 .vectors_for_model(&token, BACKFILL_BATCH_MODEL)
3499 .unwrap()
3500 .get_vectors(&entity_ids, "local", "entity.body")
3501 .await
3502 .unwrap();
3503 let single_entities = singles
3504 .vectors_for_model(&token, BACKFILL_BATCH_MODEL)
3505 .unwrap()
3506 .get_vectors(&entity_ids, "local", "entity.body")
3507 .await
3508 .unwrap();
3509 assert_eq!(batched_entities.len(), 257);
3510 assert_eq!(batched_entities, single_entities);
3511 let batched_notes = runtime
3512 .vectors_for_model(&token, BACKFILL_BATCH_MODEL)
3513 .unwrap()
3514 .get_vectors(¬e_ids, "local", "note.content")
3515 .await
3516 .unwrap();
3517 let single_notes = singles
3518 .vectors_for_model(&token, BACKFILL_BATCH_MODEL)
3519 .unwrap()
3520 .get_vectors(¬e_ids, "local", "note.content")
3521 .await
3522 .unwrap();
3523 assert_eq!(batched_notes.len(), 3);
3524 assert_eq!(batched_notes, single_notes);
3525 }
3526
3527 #[tokio::test]
3528 async fn backfill_failed_pages_retry_singly_without_duplicate_index_rows() {
3529 use khive_storage::types::{SqlStatement, SqlValue};
3530
3531 let token = NamespaceToken::local();
3532 let runtime = KhiveRuntime::memory().unwrap();
3533 for name in ["good one", "poison", "good two"] {
3534 runtime
3535 .create_entity(&token, "concept", None, name, Some("body"), None, vec![])
3536 .await
3537 .unwrap();
3538 runtime
3539 .create_note(&token, "observation", None, name, None, None, vec![])
3540 .await
3541 .unwrap();
3542 }
3543 let calls = Arc::new(AtomicUsize::new(0));
3544 runtime.register_embedder(CountingBackfillProvider {
3545 calls: Arc::clone(&calls),
3546 reject_poison: true,
3547 });
3548 let backfilled = runtime.backfill_missing_embeddings(&token).await.unwrap();
3549 assert_eq!(backfilled, 4);
3550 assert_eq!(
3551 calls.load(Ordering::SeqCst),
3552 8,
3553 "two failed pages plus six singleton retries"
3554 );
3555 let mut reader = runtime.sql().reader().await.unwrap();
3556 for (table, expected) in [("ann_write_log", 4), ("vector_provenance", 0)] {
3557 let count = reader
3558 .query_scalar(SqlStatement {
3559 sql: format!("SELECT COUNT(*) FROM {table} WHERE namespace = ?1"),
3560 params: vec![SqlValue::Text("local".into())],
3561 label: Some("backfill-batch-no-duplicate-rows".into()),
3562 })
3563 .await
3564 .unwrap();
3565 assert!(
3566 matches!(count, Some(SqlValue::Integer(value)) if value == expected),
3567 "{table} must match the guarded per-record writer"
3568 );
3569 }
3570 assert_eq!(
3571 runtime.backfill_missing_embeddings(&token).await.unwrap(),
3572 0
3573 );
3574 }
3575}