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)]
1613#[path = "retrieval_tests.rs"]
1614mod tests;