1use std::collections::{HashMap, HashSet};
4
5use uuid::Uuid;
6
7use crate::config::{parse_embedding_model_alias, sanitize_key};
8use crate::curation::note_fts_document;
9use crate::error::{RuntimeError, RuntimeResult};
10use crate::runtime::{KhiveRuntime, NamespaceToken};
11use khive_score::{rrf_score, DeterministicScore};
12use khive_storage::types::{
13 PageRequest, TextFilter, TextQueryMode, TextSearchHit, TextSearchRequest, VectorSearchHit,
14 VectorSearchRequest,
15};
16use khive_storage::EntityFilter;
17use khive_types::SubstrateKind;
18
19#[cfg(any(test, feature = "fault-injection"))]
24std::thread_local! {
25 static BACKFILL_READER_FAIL: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
26}
27
28#[cfg(any(test, feature = "fault-injection"))]
35pub fn arm_backfill_reader_fail() {
36 BACKFILL_READER_FAIL.with(|c| c.set(true));
37}
38
39#[derive(Clone, Debug)]
41pub struct SearchHit {
42 pub entity_id: Uuid,
43 pub score: DeterministicScore,
44 pub source: SearchSource,
45 pub title: Option<String>,
46 pub snippet: Option<String>,
47}
48
49#[derive(Clone, Copy, Debug, PartialEq, Eq)]
51pub enum SearchSource {
52 Vector,
53 Text,
54 Both,
55}
56
57const RRF_K: usize = 10;
65
66const CANDIDATE_MULTIPLIER: u32 = 4;
68
69impl KhiveRuntime {
70 pub async fn embed(&self, text: &str) -> RuntimeResult<Vec<f32>> {
75 let model_name = self.default_embedder_name();
76 if model_name.is_empty() {
77 return Err(RuntimeError::Unconfigured("embedding_model".into()));
78 }
79 self.embed_with_model(model_name, text).await
80 }
81
82 pub async fn embed_with_model(&self, model_name: &str, text: &str) -> RuntimeResult<Vec<f32>> {
98 let model = parse_embedding_model_alias(model_name);
102 let service = self.embedder(model_name).await?;
103 let emb_model = model.unwrap_or_default();
104 Ok(service.embed_one(text, emb_model).await?)
105 }
106
107 pub async fn embed_document_with_model(
126 &self,
127 model_name: &str,
128 text: &str,
129 ) -> RuntimeResult<Vec<f32>> {
130 let model = parse_embedding_model_alias(model_name);
131 let service = self.embedder(model_name).await?;
132 let emb_model = model.unwrap_or_default();
133 service
134 .embed_passage(&[text.to_string()], emb_model)
135 .await?
136 .into_iter()
137 .next()
138 .ok_or_else(|| RuntimeError::Internal("embed_passage returned empty vec".into()))
139 }
140
141 pub async fn embed_query_with_model(
154 &self,
155 model_name: &str,
156 text: &str,
157 ) -> RuntimeResult<Vec<f32>> {
158 let model = parse_embedding_model_alias(model_name);
159 let service = self.embedder(model_name).await?;
160 let emb_model = model.unwrap_or_default();
161 service
162 .embed_query(&[text.to_string()], emb_model)
163 .await?
164 .into_iter()
165 .next()
166 .ok_or_else(|| RuntimeError::Internal("embed_query returned empty vec".into()))
167 }
168
169 pub async fn embed_document(&self, text: &str) -> RuntimeResult<Vec<f32>> {
176 let model_name = self.default_embedder_name();
177 if model_name.is_empty() {
178 return Err(RuntimeError::Unconfigured("embedding_model".into()));
179 }
180 self.embed_document_with_model(model_name, text).await
181 }
182
183 pub async fn embed_query(&self, text: &str) -> RuntimeResult<Vec<f32>> {
190 let model_name = self.default_embedder_name();
191 if model_name.is_empty() {
192 return Err(RuntimeError::Unconfigured("embedding_model".into()));
193 }
194 self.embed_query_with_model(model_name, text).await
195 }
196
197 pub async fn embed_batch(&self, texts: &[String]) -> RuntimeResult<Vec<Vec<f32>>> {
205 if texts.is_empty() {
206 return Ok(vec![]);
207 }
208 let model_name = self.default_embedder_name();
209 if model_name.is_empty() {
210 return Err(RuntimeError::Unconfigured("embedding_model".into()));
211 }
212 self.embed_batch_with_model(model_name, texts).await
213 }
214
215 pub async fn embed_batch_with_model(
220 &self,
221 model_name: &str,
222 texts: &[String],
223 ) -> RuntimeResult<Vec<Vec<f32>>> {
224 if texts.is_empty() {
225 return Ok(vec![]);
226 }
227 let model = parse_embedding_model_alias(model_name);
228 let service = self.embedder(model_name).await?;
229 let emb_model = model.unwrap_or_default();
230 Ok(service.embed(texts, emb_model).await?)
231 }
232
233 pub async fn embed_document_batch_with_model(
243 &self,
244 model_name: &str,
245 texts: &[String],
246 ) -> RuntimeResult<Vec<Vec<f32>>> {
247 if texts.is_empty() {
248 return Ok(vec![]);
249 }
250 let model = parse_embedding_model_alias(model_name);
251 let service = self.embedder(model_name).await?;
252 let emb_model = model.unwrap_or_default();
253 Ok(service.embed_passage(texts, emb_model).await?)
254 }
255
256 pub async fn embed_document_batch(&self, texts: &[String]) -> RuntimeResult<Vec<Vec<f32>>> {
263 if texts.is_empty() {
264 return Ok(vec![]);
265 }
266 let model_name = self.default_embedder_name();
267 if model_name.is_empty() {
268 return Err(RuntimeError::Unconfigured("embedding_model".into()));
269 }
270 self.embed_document_batch_with_model(model_name, texts)
271 .await
272 }
273
274 pub async fn embed_query_batch_with_model(
281 &self,
282 model_name: &str,
283 texts: &[String],
284 ) -> RuntimeResult<Vec<Vec<f32>>> {
285 if texts.is_empty() {
286 return Ok(vec![]);
287 }
288 let model = parse_embedding_model_alias(model_name);
289 let service = self.embedder(model_name).await?;
290 let emb_model = model.unwrap_or_default();
291 Ok(service.embed_query(texts, emb_model).await?)
292 }
293
294 pub async fn vector_search(
300 &self,
301 token: &NamespaceToken,
302 query_embedding: Option<Vec<f32>>,
303 query_text: Option<&str>,
304 top_k: u32,
305 kind: Option<SubstrateKind>,
306 ) -> RuntimeResult<Vec<VectorSearchHit>> {
307 let embedding = match query_embedding {
308 Some(vec) => vec,
309 None => {
310 let text = query_text.ok_or_else(|| {
311 RuntimeError::InvalidInput(
312 "vector search requires query_embedding or query_text".into(),
313 )
314 })?;
315 if text.trim().is_empty() {
316 return Err(RuntimeError::InvalidInput(
317 "query_text must not be empty".into(),
318 ));
319 }
320 self.embed_query(text).await?
321 }
322 };
323
324 let ns = token.namespace().as_str().to_owned();
325 Ok(self
326 .vectors(token)?
327 .search(VectorSearchRequest {
328 query_vectors: vec![embedding],
329 top_k,
330 namespace: Some(ns),
331 kind,
332 embedding_model: None,
333 filter: None,
334 backend_hints: None,
335 })
336 .await?)
337 }
338
339 #[allow(clippy::too_many_arguments)]
384 pub async fn hybrid_search(
385 &self,
386 token: &NamespaceToken,
387 query_text: &str,
388 query_vector: Option<Vec<f32>>,
389 limit: u32,
390 entity_kind: Option<&str>,
391 entity_type: Option<&str>,
392 tags_any: &[String],
393 properties_filter: Option<&serde_json::Value>,
394 ) -> RuntimeResult<Vec<SearchHit>> {
395 let candidates = limit.saturating_mul(CANDIDATE_MULTIPLIER).max(limit);
396
397 let visible_ns: Vec<String> = token
398 .visible_namespaces()
399 .iter()
400 .map(|ns| ns.as_str().to_owned())
401 .collect();
402 let text_hits = self
403 .text(token)?
404 .search(TextSearchRequest {
405 query: query_text.to_string(),
406 mode: TextQueryMode::Plain,
407 filter: Some(TextFilter {
408 namespaces: visible_ns.clone(),
409 ..TextFilter::default()
410 }),
411 top_k: candidates,
412 snippet_chars: 200,
413 })
414 .await?;
415
416 let vector_hits = if query_vector.is_some() || self.config().embedding_model.is_some() {
417 self.vector_search(
418 token,
419 query_vector,
420 Some(query_text),
421 candidates,
422 Some(SubstrateKind::Entity),
423 )
424 .await?
425 } else {
426 Vec::new()
427 };
428
429 let mut fused = rrf_fuse(text_hits, vector_hits, candidates as usize, query_text);
433
434 if !fused.is_empty() {
441 let candidate_ids: Vec<Uuid> = fused.iter().map(|h| h.entity_id).collect();
442 let alive_page = self
443 .entities(token)?
444 .query_entities(
445 token.namespace().as_str(),
446 EntityFilter {
447 ids: candidate_ids,
448 kinds: entity_kind.map(|k| vec![k.to_string()]).unwrap_or_default(),
449 entity_types: entity_type.map(|t| vec![t.to_string()]).unwrap_or_default(),
450 namespaces: visible_ns,
451 tags_any: tags_any.to_vec(),
452 ..EntityFilter::default()
453 },
454 PageRequest {
455 offset: 0,
456 limit: fused.len() as u32,
457 },
458 )
459 .await?;
460 let mut entity_meta: HashMap<Uuid, (String, Option<String>)> = HashMap::new();
463 let mut alive: HashSet<Uuid> = HashSet::new();
464 for e in alive_page.items {
465 if let Some(pf) = properties_filter {
468 if !entity_props_match(e.properties.as_ref(), pf) {
469 continue;
470 }
471 }
472 alive.insert(e.id);
473 entity_meta.insert(e.id, (e.name, e.description));
474 }
475
476 fused.retain(|h| alive.contains(&h.entity_id));
477
478 for hit in &mut fused {
480 if let Some((name, description)) = entity_meta.get(&hit.entity_id) {
481 if hit.title.is_none() {
482 hit.title = Some(name.clone());
483 }
484 if hit.snippet.is_none() {
485 hit.snippet = description.clone();
486 }
487 }
488 }
489 }
490
491 fused.truncate(limit as usize);
492 Ok(fused)
493 }
494
495 pub async fn knn(
501 &self,
502 token: &NamespaceToken,
503 query_vector: Vec<f32>,
504 top_k: u32,
505 ) -> RuntimeResult<Vec<VectorSearchHit>> {
506 let ns = token.namespace().as_str().to_owned();
507 Ok(self
508 .vectors(token)?
509 .search(VectorSearchRequest {
510 query_vectors: vec![query_vector],
511 top_k,
512 namespace: Some(ns),
513 kind: Some(SubstrateKind::Entity),
514 embedding_model: None,
515 filter: None,
516 backend_hints: None,
517 })
518 .await?)
519 }
520
521 pub async fn rerank(
527 &self,
528 token: &NamespaceToken,
529 query_vector: &[f32],
530 candidate_ids: &[Uuid],
531 top_k: u32,
532 ) -> RuntimeResult<Vec<VectorSearchHit>> {
533 let candidate_set: HashSet<Uuid> = candidate_ids.iter().copied().collect();
534 let ns = token.namespace().as_str().to_owned();
535 let all_hits = self
536 .vectors(token)?
537 .search(VectorSearchRequest {
538 query_vectors: vec![query_vector.to_vec()],
539 top_k: candidate_ids.len() as u32,
540 namespace: Some(ns),
541 kind: Some(SubstrateKind::Entity),
542 embedding_model: None,
543 filter: None,
544 backend_hints: None,
545 })
546 .await?;
547 let mut hits: Vec<VectorSearchHit> = all_hits
548 .into_iter()
549 .filter(|h| candidate_set.contains(&h.subject_id))
550 .collect();
551 hits.sort_by(|a, b| b.score.cmp(&a.score));
552 hits.truncate(top_k as usize);
553 Ok(hits)
554 }
555
556 pub async fn backfill_missing_embeddings(&self, token: &NamespaceToken) -> RuntimeResult<u64> {
569 use khive_storage::types::{SqlRow, SqlStatement, SqlValue};
570
571 let model_names = self.registered_embedding_model_names();
572 if model_names.is_empty() {
573 tracing::debug!(
574 "backfill_missing_embeddings: no embedding models registered, skipping"
575 );
576 return Ok(0);
577 }
578
579 let ns = token.namespace().as_str().to_string();
580 let mut total_backfilled = 0u64;
581
582 for model_name in &model_names {
583 let vec_table = format!("vec_{}", sanitize_key(model_name));
585
586 const PAGE_SIZE: usize = 500;
591 let mut entity_total = 0usize;
592 loop {
593 let entity_sql = SqlStatement {
594 sql: format!(
595 "SELECT id, name, description FROM entities \
596 WHERE namespace = ?1 AND deleted_at IS NULL \
597 AND id NOT IN (\
598 SELECT subject_id FROM {vec_table} \
599 WHERE namespace = ?1 AND embedding_model = ?2 \
600 ) LIMIT {PAGE_SIZE}"
601 ),
602 params: vec![
603 SqlValue::Text(ns.clone()),
604 SqlValue::Text(model_name.clone()),
605 ],
606 label: Some("backfill_entities".into()),
607 };
608
609 let entity_rows: Vec<SqlRow> = {
610 let sql = self.sql();
611 let reader_result = sql.reader().await;
612 #[cfg(any(test, feature = "fault-injection"))]
613 let reader_result = if BACKFILL_READER_FAIL.with(|c| c.get()) {
614 BACKFILL_READER_FAIL.with(|c| c.set(false));
615 Err(khive_storage::StorageError::Pool {
616 operation: "reader".into(),
617 message: "injected failure".into(),
618 })
619 } else {
620 reader_result
621 };
622 let mut reader = reader_result.map_err(RuntimeError::Storage)?;
623 reader
624 .query_all(entity_sql)
625 .await
626 .map_err(RuntimeError::Storage)?
627 };
628
629 let batch_len = entity_rows.len();
630 entity_total += batch_len;
631
632 for row in &entity_rows {
633 let id_str = row.columns.first().and_then(|c| {
634 if let SqlValue::Text(s) = &c.value {
635 Some(s.clone())
636 } else {
637 None
638 }
639 });
640 let description = row.columns.get(2).and_then(|c| {
641 if let SqlValue::Text(s) = &c.value {
642 Some(s.clone())
643 } else if let SqlValue::Null = &c.value {
644 None
645 } else {
646 None
647 }
648 });
649
650 let (Some(id_str), Some(desc)) = (id_str, description) else {
651 continue;
652 };
653 let Ok(id) = id_str.parse::<Uuid>() else {
654 continue;
655 };
656 if desc.trim().is_empty() {
657 continue;
658 }
659
660 match self.embed_document_with_model(model_name, &desc).await {
661 Ok(vector) => {
662 if let Ok(vs) = self.vectors_for_model(token, model_name) {
663 match vs
664 .insert(
665 id,
666 SubstrateKind::Entity,
667 &ns,
668 "entity.description",
669 vec![vector],
670 )
671 .await
672 {
673 Ok(()) => {
674 total_backfilled += 1;
675 }
676 Err(e) => {
677 tracing::warn!(
678 id = %id, model = %model_name,
679 error = %e,
680 "backfill_missing_embeddings: entity vector insert failed"
681 );
682 }
683 }
684 }
685 }
686 Err(e) => {
687 tracing::warn!(
688 id = %id, model = %model_name,
689 error = %e,
690 "backfill_missing_embeddings: entity embed failed"
691 );
692 }
693 }
694 }
695
696 if batch_len < PAGE_SIZE {
697 break;
698 }
699 }
700
701 let text_store = self.text_for_notes(token).ok();
703 let note_store = self.notes(token).ok();
704 let mut note_total = 0usize;
705 loop {
706 let note_sql = SqlStatement {
710 sql: format!(
711 "SELECT id FROM notes \
712 WHERE namespace = ?1 AND deleted_at IS NULL \
713 AND id NOT IN (\
714 SELECT subject_id FROM {vec_table} \
715 WHERE namespace = ?1 AND embedding_model = ?2 \
716 ) LIMIT {PAGE_SIZE}"
717 ),
718 params: vec![
719 SqlValue::Text(ns.clone()),
720 SqlValue::Text(model_name.clone()),
721 ],
722 label: Some("backfill_notes".into()),
723 };
724
725 let note_rows: Vec<SqlRow> = {
726 let sql = self.sql();
727 let reader_result = sql.reader().await;
728 #[cfg(any(test, feature = "fault-injection"))]
729 let reader_result = if BACKFILL_READER_FAIL.with(|c| c.get()) {
730 BACKFILL_READER_FAIL.with(|c| c.set(false));
731 Err(khive_storage::StorageError::Pool {
732 operation: "reader".into(),
733 message: "injected failure".into(),
734 })
735 } else {
736 reader_result
737 };
738 let mut reader = reader_result.map_err(RuntimeError::Storage)?;
739 reader
740 .query_all(note_sql)
741 .await
742 .map_err(RuntimeError::Storage)?
743 };
744
745 let batch_len = note_rows.len();
746 note_total += batch_len;
747
748 for row in ¬e_rows {
749 let id_str = row.columns.first().and_then(|c| {
750 if let SqlValue::Text(s) = &c.value {
751 Some(s.clone())
752 } else {
753 None
754 }
755 });
756
757 let Some(id_str) = id_str else {
758 continue;
759 };
760 let Ok(id) = id_str.parse::<Uuid>() else {
761 continue;
762 };
763
764 let note = match ¬e_store {
768 Some(store) => match store.get_note(id).await {
769 Ok(Some(n)) => n,
770 _ => continue,
771 },
772 None => continue,
773 };
774
775 if note.content.trim().is_empty() {
776 continue;
777 }
778
779 if model_names.first().map(|n| n.as_str()) == Some(model_name.as_str()) {
782 if let Some(ref ts) = text_store {
783 if let Err(e) = ts.upsert_document(note_fts_document(¬e)).await {
784 tracing::warn!(id = %id, error = %e,
785 "backfill_missing_embeddings: note FTS upsert failed");
786 }
787 }
788 }
789
790 let content = note.content.clone();
791 match self.embed_document_with_model(model_name, &content).await {
792 Ok(vector) => {
793 if let Ok(vs) = self.vectors_for_model(token, model_name) {
794 match vs
795 .insert(
796 id,
797 SubstrateKind::Note,
798 &ns,
799 "note.content",
800 vec![vector],
801 )
802 .await
803 {
804 Ok(()) => {
805 total_backfilled += 1;
806 }
807 Err(e) => {
808 tracing::warn!(
809 id = %id, model = %model_name,
810 error = %e,
811 "backfill_missing_embeddings: note vector insert failed"
812 );
813 }
814 }
815 }
816 }
817 Err(e) => {
818 tracing::warn!(
819 id = %id, model = %model_name,
820 error = %e,
821 "backfill_missing_embeddings: note embed failed"
822 );
823 }
824 }
825 }
826
827 if batch_len < PAGE_SIZE {
828 break;
829 }
830 }
831
832 tracing::info!(
833 model = %model_name,
834 namespace = %ns,
835 entities = entity_total,
836 notes = note_total,
837 "backfill_missing_embeddings: model pass complete"
838 );
839 }
840
841 tracing::info!(
842 namespace = %ns,
843 total_backfilled = total_backfilled,
844 "backfill_missing_embeddings: finished"
845 );
846
847 Ok(total_backfilled)
848 }
849
850 pub async fn sweep_orphan_vectors(
865 &self,
866 token: &NamespaceToken,
867 max_delete_per_model: u32,
868 dry_run: bool,
869 ) -> RuntimeResult<u64> {
870 use khive_storage::types::OrphanSweepConfig;
871 use khive_storage::StorageError;
872
873 let model_names = self.registered_embedding_model_names();
874 if model_names.is_empty() {
875 tracing::debug!("sweep_orphan_vectors: no embedding models registered, skipping");
876 return Ok(0);
877 }
878
879 let ns = token.namespace().as_str().to_string();
880 let mut total_deleted = 0u64;
881
882 for model_name in &model_names {
883 let store = match self.vectors_for_model(token, model_name) {
884 Ok(s) => s,
885 Err(e) => {
886 tracing::warn!(
887 model = %model_name,
888 error = %e,
889 "sweep_orphan_vectors: failed to get vector store, skipping model"
890 );
891 continue;
892 }
893 };
894
895 let caps = store.capabilities();
896 if !caps.supports_orphan_sweep {
897 tracing::debug!(
898 model = %model_name,
899 "sweep_orphan_vectors: backend does not support orphan sweep, skipping"
900 );
901 continue;
902 }
903
904 let config = OrphanSweepConfig {
905 subject_id_allowlist: None,
906 namespaces: vec![ns.clone()],
907 substrate_kinds: vec![],
908 max_delete: max_delete_per_model,
909 dry_run,
910 };
911
912 match store.orphan_sweep(&config).await {
913 Ok(result) => {
914 tracing::info!(
915 model = %model_name,
916 namespace = %ns,
917 scanned = result.scanned,
918 deleted = result.deleted,
919 would_delete = result.would_delete,
920 dry_run = dry_run,
921 "sweep_orphan_vectors: sweep complete"
922 );
923 total_deleted += result.deleted;
924 }
925 Err(StorageError::Unsupported { .. }) => {
926 tracing::debug!(
927 model = %model_name,
928 "sweep_orphan_vectors: backend returned Unsupported, skipping"
929 );
930 }
931 Err(e) => {
932 tracing::warn!(
933 model = %model_name,
934 error = %e,
935 "sweep_orphan_vectors: sweep failed, continuing with other models"
936 );
937 }
938 }
939 }
940
941 tracing::info!(
942 namespace = %ns,
943 total_deleted = total_deleted,
944 dry_run = dry_run,
945 "sweep_orphan_vectors: finished"
946 );
947
948 Ok(total_deleted)
949 }
950}
951
952fn entity_props_match(
957 entity_props: Option<&serde_json::Value>,
958 filter: &serde_json::Value,
959) -> bool {
960 let required = match filter.as_object() {
961 Some(obj) if !obj.is_empty() => obj,
962 _ => return true,
963 };
964 let actual = match entity_props.and_then(serde_json::Value::as_object) {
965 Some(obj) => obj,
966 None => return false,
967 };
968 required
969 .iter()
970 .all(|(k, v)| actual.get(k).is_some_and(|av| av == v))
971}
972
973const EXACT_MATCH_BOOST: f64 = 0.5;
977
978fn rrf_fuse(
986 text_hits: Vec<TextSearchHit>,
987 vector_hits: Vec<VectorSearchHit>,
988 limit: usize,
989 query_text: &str,
990) -> Vec<SearchHit> {
991 #[derive(Default)]
992 struct Bucket {
993 score: DeterministicScore,
994 source: Option<SearchSource>,
995 title: Option<String>,
996 snippet: Option<String>,
997 }
998
999 let mut buckets: HashMap<Uuid, Bucket> = HashMap::new();
1000
1001 let query_lower = query_text.to_lowercase();
1002 for (i, hit) in text_hits.into_iter().enumerate() {
1003 let rank = i + 1; let entry = buckets.entry(hit.subject_id).or_default();
1005 entry.score = entry.score + rrf_score(rank, RRF_K);
1006 entry.source = Some(match entry.source {
1007 Some(SearchSource::Vector) => SearchSource::Both,
1008 _ => SearchSource::Text,
1009 });
1010 if entry.title.is_none() {
1011 if let Some(ref title) = hit.title {
1013 if title.to_lowercase() == query_lower {
1014 entry.score = entry.score + DeterministicScore::from_f64(EXACT_MATCH_BOOST);
1015 }
1016 }
1017 entry.title = hit.title;
1018 }
1019 if entry.snippet.is_none() {
1020 entry.snippet = hit.snippet;
1021 }
1022 }
1023
1024 for (i, hit) in vector_hits.into_iter().enumerate() {
1025 let rank = i + 1;
1026 let entry = buckets.entry(hit.subject_id).or_default();
1027 entry.score = entry.score + rrf_score(rank, RRF_K);
1028 entry.source = Some(match entry.source {
1029 Some(SearchSource::Text) => SearchSource::Both,
1030 _ => SearchSource::Vector,
1031 });
1032 }
1033
1034 let mut hits: Vec<SearchHit> = buckets
1035 .into_iter()
1036 .map(|(id, b)| SearchHit {
1037 entity_id: id,
1038 score: b.score,
1039 source: b.source.expect("each bucket gets a source"),
1040 title: b.title,
1041 snippet: b.snippet,
1042 })
1043 .collect();
1044
1045 hits.sort_by(|a, b| b.score.cmp(&a.score).then(a.entity_id.cmp(&b.entity_id)));
1046 hits.truncate(limit);
1047 hits
1048}
1049
1050#[cfg(test)]
1051mod tests {
1052 use super::*;
1053 use crate::runtime::{KhiveRuntime, NamespaceToken, RuntimeConfig};
1054 use khive_storage::types::{TextSearchHit, VectorSearchHit};
1055 use khive_types::namespace::Namespace;
1056 use lattice_embed::EmbeddingModel;
1057
1058 fn text_hit(id: Uuid, rank: u32, title: &str) -> TextSearchHit {
1059 TextSearchHit {
1060 subject_id: id,
1061 score: DeterministicScore::from_f64(1.0),
1062 rank,
1063 title: Some(title.to_string()),
1064 snippet: Some("...".to_string()),
1065 }
1066 }
1067
1068 fn vector_hit(id: Uuid, rank: u32) -> VectorSearchHit {
1069 VectorSearchHit {
1070 subject_id: id,
1071 score: DeterministicScore::from_f64(0.9),
1072 rank,
1073 }
1074 }
1075
1076 #[test]
1077 fn rrf_fuse_text_only() {
1078 let a = Uuid::new_v4();
1079 let b = Uuid::new_v4();
1080 let text = vec![text_hit(a, 1, "A"), text_hit(b, 2, "B")];
1081 let hits = rrf_fuse(text, vec![], 10, "query");
1082 assert_eq!(hits.len(), 2);
1083 assert_eq!(hits[0].entity_id, a);
1084 assert_eq!(hits[0].source, SearchSource::Text);
1085 assert_eq!(hits[0].title.as_deref(), Some("A"));
1086 }
1087
1088 #[test]
1089 fn rrf_fuse_vector_only() {
1090 let a = Uuid::new_v4();
1091 let hits = rrf_fuse(vec![], vec![vector_hit(a, 1)], 10, "query");
1092 assert_eq!(hits.len(), 1);
1093 assert_eq!(hits[0].source, SearchSource::Vector);
1094 assert!(hits[0].title.is_none());
1095 }
1096
1097 #[test]
1098 fn rrf_fuse_marks_both_when_in_both_lists() {
1099 let id = Uuid::new_v4();
1100 let text = vec![text_hit(id, 1, "A")];
1101 let vec = vec![vector_hit(id, 1)];
1102 let hits = rrf_fuse(text, vec, 10, "query");
1103 assert_eq!(hits.len(), 1);
1104 assert_eq!(hits[0].source, SearchSource::Both);
1105 }
1106
1107 #[test]
1108 fn rrf_fuse_respects_limit() {
1109 let hits: Vec<TextSearchHit> = (0..20)
1110 .map(|i| text_hit(Uuid::new_v4(), i + 1, "x"))
1111 .collect();
1112 let fused = rrf_fuse(hits, vec![], 5, "query");
1113 assert_eq!(fused.len(), 5);
1114 }
1115
1116 #[test]
1117 fn rrf_fuse_orders_higher_score_first() {
1118 let a = Uuid::new_v4();
1120 let b = Uuid::new_v4();
1121 let text = vec![text_hit(a, 1, "A")];
1122 let vec = vec![vector_hit(a, 1), vector_hit(b, 2)];
1123 let hits = rrf_fuse(text, vec, 10, "query");
1124 assert_eq!(hits[0].entity_id, a);
1125 assert_eq!(hits[0].source, SearchSource::Both);
1126 assert!(hits[0].score > hits[1].score);
1127 }
1128
1129 #[test]
1130 fn rrf_fuse_k10_score_spread_exceeds_threshold() {
1131 let ids: Vec<Uuid> = (0..10).map(|_| Uuid::new_v4()).collect();
1134 let text: Vec<TextSearchHit> = ids
1135 .iter()
1136 .enumerate()
1137 .map(|(i, &id)| text_hit(id, (i + 1) as u32, "x"))
1138 .collect();
1139 let hits = rrf_fuse(text, vec![], 10, "query");
1140 assert_eq!(hits.len(), 10);
1141 let top_score = hits[0].score.to_f64();
1142 let bottom_score = hits[9].score.to_f64();
1143 let spread = top_score - bottom_score;
1144 assert!(
1145 spread >= 0.03,
1146 "score spread {spread:.4} between rank 1 and rank 10 must be ≥ 0.03 (was {spread:.4})"
1147 );
1148 }
1149
1150 #[test]
1151 fn rrf_fuse_exact_match_boost_elevates_score() {
1152 let exact_id = Uuid::new_v4();
1155 let other_id = Uuid::new_v4();
1156 let text = vec![
1158 text_hit(other_id, 1, "something else"),
1159 text_hit(exact_id, 2, "FlashAttention"),
1160 ];
1161 let hits = rrf_fuse(text, vec![], 10, "flashattention");
1162 assert_eq!(hits.len(), 2);
1163 assert_eq!(
1164 hits[0].entity_id, exact_id,
1165 "exact match must rank first despite being rank-2 in raw text search"
1166 );
1167 }
1168
1169 #[test]
1172 fn embed_batch_unconfigured_on_memory_runtime() {
1173 let rt = KhiveRuntime::memory().unwrap();
1175 let result = tokio::runtime::Runtime::new()
1176 .unwrap()
1177 .block_on(rt.embed_batch(&[]));
1178 assert!(result.is_ok());
1180 assert!(result.unwrap().is_empty());
1181 }
1182
1183 #[test]
1184 fn embed_batch_empty_input_returns_empty_vec() {
1185 let rt = KhiveRuntime::memory().unwrap();
1187 let result = tokio::runtime::Runtime::new()
1188 .unwrap()
1189 .block_on(rt.embed_batch(&[]));
1190 assert_eq!(result.unwrap(), Vec::<Vec<f32>>::new());
1191 }
1192
1193 #[test]
1194 fn embed_batch_no_model_non_empty_returns_unconfigured() {
1195 let rt = KhiveRuntime::memory().unwrap();
1196 let texts = vec!["hello".to_string()];
1197 let result = tokio::runtime::Runtime::new()
1198 .unwrap()
1199 .block_on(rt.embed_batch(&texts));
1200 match result {
1201 Err(crate::RuntimeError::Unconfigured(s)) => assert_eq!(s, "embedding_model"),
1202 Err(other) => panic!("expected Unconfigured, got {:?}", other),
1203 Ok(_) => panic!("expected Err, got Ok"),
1204 }
1205 }
1206
1207 #[test]
1208 #[ignore = "loads ~80 MB model; run with --include-ignored"]
1209 fn embed_batch_count_matches_input() {
1210 let config = RuntimeConfig {
1211 db_path: None,
1212 default_namespace: Namespace::parse("test").unwrap(),
1213 embedding_model: Some(EmbeddingModel::AllMiniLmL6V2),
1214 packs: vec!["kg".to_string()],
1215 ..RuntimeConfig::default()
1216 };
1217 let rt = KhiveRuntime::new(config).unwrap();
1218 let texts: Vec<String> = vec!["foo".to_string(), "bar".to_string(), "baz".to_string()];
1219 let result = tokio::runtime::Runtime::new()
1220 .unwrap()
1221 .block_on(rt.embed_batch(&texts));
1222 let embeddings = result.unwrap();
1223 assert_eq!(embeddings.len(), texts.len());
1224 }
1225
1226 #[test]
1227 fn vector_search_requires_embedding_or_text() {
1228 let rt = KhiveRuntime::memory().unwrap();
1229 let tok = NamespaceToken::local();
1230 let result = tokio::runtime::Runtime::new()
1231 .unwrap()
1232 .block_on(rt.vector_search(&tok, None, None, 10, Some(SubstrateKind::Entity)));
1233 match result {
1234 Err(crate::RuntimeError::InvalidInput(msg)) => {
1235 assert!(msg.contains("query_embedding or query_text"), "msg: {msg}");
1236 }
1237 other => panic!("expected InvalidInput, got {other:?}"),
1238 }
1239 }
1240
1241 #[test]
1242 fn vector_search_text_without_model_returns_unconfigured() {
1243 let rt = KhiveRuntime::memory().unwrap();
1244 let tok = NamespaceToken::local();
1245 let result = tokio::runtime::Runtime::new()
1246 .unwrap()
1247 .block_on(rt.vector_search(
1248 &tok,
1249 None,
1250 Some("attention"),
1251 10,
1252 Some(SubstrateKind::Entity),
1253 ));
1254 match result {
1255 Err(crate::RuntimeError::Unconfigured(s)) => assert_eq!(s, "embedding_model"),
1256 other => panic!("expected Unconfigured, got {other:?}"),
1257 }
1258 }
1259
1260 #[test]
1261 #[ignore = "loads ~80 MB model; run with --include-ignored"]
1262 fn embed_batch_vectors_have_expected_dimensions() {
1263 let model = EmbeddingModel::AllMiniLmL6V2;
1264 let config = RuntimeConfig {
1265 db_path: None,
1266 default_namespace: Namespace::parse("test").unwrap(),
1267 embedding_model: Some(model),
1268 packs: vec!["kg".to_string()],
1269 ..RuntimeConfig::default()
1270 };
1271 let rt = KhiveRuntime::new(config).unwrap();
1272 let texts = vec!["hello world".to_string()];
1273 let result = tokio::runtime::Runtime::new()
1274 .unwrap()
1275 .block_on(rt.embed_batch(&texts));
1276 let embeddings = result.unwrap();
1277 assert_eq!(embeddings[0].len(), model.dimensions());
1278 }
1279
1280 #[tokio::test]
1283 async fn hybrid_search_entity_hit_has_title() {
1284 let rt = KhiveRuntime::memory().unwrap();
1285 let tok = NamespaceToken::local();
1286 rt.create_entity(
1287 &tok,
1288 "concept",
1289 None,
1290 "FlashAttention",
1291 Some("IO-aware exact attention using tiling"),
1292 None,
1293 vec![],
1294 )
1295 .await
1296 .unwrap();
1297
1298 let hits = rt
1299 .hybrid_search(&tok, "FlashAttention", None, 10, None, None, &[], None)
1300 .await
1301 .unwrap();
1302
1303 assert!(!hits.is_empty(), "should find the entity");
1304 let hit = &hits[0];
1305 assert!(hit.title.is_some(), "title must be populated");
1306 assert!(
1307 hit.title.as_deref().unwrap().contains("FlashAttention"),
1308 "title must contain entity name"
1309 );
1310 }
1311
1312 #[tokio::test]
1331 async fn hybrid_search_tag_filter_pushed_before_truncation() {
1332 let rt = KhiveRuntime::memory().unwrap();
1333 let tok = NamespaceToken::local();
1334
1335 rt.create_entity(
1337 &tok,
1338 "concept",
1339 None,
1340 "alpha beta gamma decoy alpha beta gamma",
1341 Some("alpha beta gamma decoy description alpha beta gamma"),
1342 None,
1343 vec!["other-tag".to_string()],
1344 )
1345 .await
1346 .unwrap();
1347
1348 let target = rt
1350 .create_entity(
1351 &tok,
1352 "concept",
1353 None,
1354 "alpha beta gamma target",
1355 Some("alpha beta gamma target description"),
1356 None,
1357 vec!["target-tag".to_string()],
1358 )
1359 .await
1360 .unwrap();
1361
1362 let hits = rt
1366 .hybrid_search(
1367 &tok,
1368 "alpha beta gamma",
1369 None,
1370 1,
1371 None,
1372 None,
1373 &["target-tag".to_string()],
1374 None,
1375 )
1376 .await
1377 .unwrap();
1378
1379 assert_eq!(
1380 hits.len(),
1381 1,
1382 "exactly one hit expected (the tag-matching entity)"
1383 );
1384 assert_eq!(
1385 hits[0].entity_id, target.id,
1386 "the tag-filtered entity must be returned even when ranked below limit in raw fusion"
1387 );
1388 }
1389
1390 #[tokio::test]
1399 async fn hybrid_search_props_filter_pushed_before_truncation() {
1400 let rt = KhiveRuntime::memory().unwrap();
1401 let tok = NamespaceToken::local();
1402
1403 rt.create_entity(
1404 &tok,
1405 "concept",
1406 None,
1407 "delta epsilon zeta decoy delta epsilon zeta",
1408 Some("delta epsilon zeta decoy description delta epsilon zeta"),
1409 Some(serde_json::json!({"domain": "other"})),
1410 vec![],
1411 )
1412 .await
1413 .unwrap();
1414
1415 let target = rt
1416 .create_entity(
1417 &tok,
1418 "concept",
1419 None,
1420 "delta epsilon zeta target",
1421 Some("delta epsilon zeta target description"),
1422 Some(serde_json::json!({"domain": "target"})),
1423 vec![],
1424 )
1425 .await
1426 .unwrap();
1427
1428 let filter = serde_json::json!({"domain": "target"});
1429 let hits = rt
1430 .hybrid_search(
1431 &tok,
1432 "delta epsilon zeta",
1433 None,
1434 1,
1435 None,
1436 None,
1437 &[],
1438 Some(&filter),
1439 )
1440 .await
1441 .unwrap();
1442
1443 assert_eq!(hits.len(), 1, "exactly one hit expected (properties match)");
1444 assert_eq!(
1445 hits[0].entity_id, target.id,
1446 "the properties-filtered entity must be returned even when ranked below limit"
1447 );
1448 }
1449
1450 #[test]
1453 #[ignore = "loads ~80 MB model; run with --include-ignored"]
1454 fn minilm_document_and_query_embed_are_identical_no_prefix_model() {
1455 let model = EmbeddingModel::AllMiniLmL6V2;
1459 let config = RuntimeConfig {
1460 db_path: None,
1461 default_namespace: Namespace::parse("test").unwrap(),
1462 embedding_model: Some(model),
1463 packs: vec!["kg".to_string()],
1464 ..RuntimeConfig::default()
1465 };
1466 let rt = KhiveRuntime::new(config).unwrap();
1467 let text = "attention is all you need".to_string();
1468 let rt_ref = &rt;
1469 let (doc_emb, query_emb) = tokio::runtime::Runtime::new().unwrap().block_on(async {
1470 let d = rt_ref
1471 .embed_document_with_model(&model.to_string(), &text)
1472 .await
1473 .unwrap();
1474 let q = rt_ref
1475 .embed_query_with_model(&model.to_string(), &text)
1476 .await
1477 .unwrap();
1478 (d, q)
1479 });
1480 assert_eq!(
1481 doc_emb, query_emb,
1482 "MiniLM has no instruction prefix: document and query embeds must be identical"
1483 );
1484 }
1485
1486 #[test]
1487 #[ignore = "loads multilingual-e5-small (~90 MB); run with --include-ignored"]
1488 fn e5_document_and_query_embed_differ_instruction_tuned_model() {
1489 let model = EmbeddingModel::MultilingualE5Small;
1494 let config = RuntimeConfig {
1495 db_path: None,
1496 default_namespace: Namespace::parse("test").unwrap(),
1497 embedding_model: Some(model),
1498 packs: vec!["kg".to_string()],
1499 ..RuntimeConfig::default()
1500 };
1501 let rt = KhiveRuntime::new(config).unwrap();
1502 let text = "attention is all you need".to_string();
1503 let rt_ref = &rt;
1504 let (doc_emb, query_emb) = tokio::runtime::Runtime::new().unwrap().block_on(async {
1505 let d = rt_ref
1506 .embed_document_with_model(&model.to_string(), &text)
1507 .await
1508 .unwrap();
1509 let q = rt_ref
1510 .embed_query_with_model(&model.to_string(), &text)
1511 .await
1512 .unwrap();
1513 (d, q)
1514 });
1515 assert_ne!(
1516 doc_emb, query_emb,
1517 "multilingual-e5-small uses asymmetric prefixes: document ('passage: ') \
1518 and query ('query: ') embeds of the same text must differ"
1519 );
1520 }
1521
1522 use crate::embedder_registry::EmbedderProvider;
1525 use lattice_embed::EmbeddingService;
1526
1527 struct StubEmbedderProvider;
1532
1533 #[async_trait::async_trait]
1534 impl EmbedderProvider for StubEmbedderProvider {
1535 fn name(&self) -> &str {
1536 "stub-model-m07"
1537 }
1538
1539 fn dimensions(&self) -> usize {
1540 4
1541 }
1542
1543 async fn build(&self) -> crate::error::RuntimeResult<std::sync::Arc<dyn EmbeddingService>> {
1544 struct StubSvc;
1545 #[async_trait::async_trait]
1546 impl EmbeddingService for StubSvc {
1547 async fn embed(
1548 &self,
1549 _texts: &[String],
1550 _model: lattice_embed::EmbeddingModel,
1551 ) -> std::result::Result<Vec<Vec<f32>>, lattice_embed::EmbedError> {
1552 Ok(vec![])
1553 }
1554
1555 fn supports_model(&self, _model: lattice_embed::EmbeddingModel) -> bool {
1556 true
1557 }
1558
1559 fn name(&self) -> &'static str {
1560 "stub-svc-m07"
1561 }
1562 }
1563 Ok(std::sync::Arc::new(StubSvc))
1564 }
1565 }
1566
1567 #[tokio::test]
1579 async fn backfill_reader_error_is_propagated_not_swallowed() {
1580 let rt = KhiveRuntime::memory().unwrap();
1581 rt.register_embedder(StubEmbedderProvider);
1582 let tok = NamespaceToken::local();
1583
1584 super::arm_backfill_reader_fail();
1587
1588 let result = rt.backfill_missing_embeddings(&tok).await;
1589 assert!(
1590 result.is_err(),
1591 "backfill_missing_embeddings must propagate the reader error (got Ok instead)"
1592 );
1593 let err_msg = result.unwrap_err().to_string();
1594 assert!(
1595 err_msg.contains("injected failure"),
1596 "error must originate from the injected reader failure, got: {err_msg}"
1597 );
1598 }
1599}