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
23const CUSTOM_QUERY_ROLE_MODEL: EmbeddingModel = EmbeddingModel::AllMiniLmL6V2;
29
30pub use khive_retrieval::{
31 HybridSearchOutcome, RankScoreKind, SearchHit, SearchSignals, SearchSource,
32};
33
34mod single_engine;
35
36#[cfg(test)]
37mod request_stop_tests;
38
39pub(crate) const EMBEDDING_BATCH_PAGE_SIZE: usize = 256;
41
42#[cfg(any(test, feature = "fault-injection"))]
44std::thread_local! {
45 static BACKFILL_READER_FAIL: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
46}
47
48#[cfg(any(test, feature = "fault-injection"))]
54pub fn arm_backfill_reader_fail() {
55 BACKFILL_READER_FAIL.with(|c| c.set(true));
56}
57
58const RRF_K: usize = 10;
65
66const CANDIDATE_MULTIPLIER: u32 = 4;
68
69pub const EMBEDDING_INPUT_TRUNCATED_WARNING: &str =
71 "embedding input was truncated to the embedder maximum; full content was stored unchanged";
72
73#[derive(Clone, Debug)]
76pub struct DocumentEmbeddingOutcome {
77 pub model_name: String,
78 pub vector: Vec<f32>,
79 pub prepared_text_fingerprint: Option<ContentRef>,
81 pub source_bytes: usize,
82 pub embedded_bytes: usize,
83 pub truncated: bool,
84}
85
86#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
88pub struct EmbeddingTruncationReport {
89 pub truncated: u64,
90 pub discarded_bytes: u64,
91}
92
93impl EmbeddingTruncationReport {
94 pub fn observe(&mut self, outcome: &DocumentEmbeddingOutcome) {
95 if outcome.truncated {
96 self.truncated += 1;
97 self.discarded_bytes +=
98 outcome.source_bytes.saturating_sub(outcome.embedded_bytes) as u64;
99 }
100 }
101
102 #[must_use]
103 pub const fn any_truncated(&self) -> bool {
104 self.truncated > 0
105 }
106
107 pub fn merge(&mut self, other: Self) {
109 self.truncated = self.truncated.saturating_add(other.truncated);
110 self.discarded_bytes = self.discarded_bytes.saturating_add(other.discarded_bytes);
111 }
112}
113
114pub fn document_embedding_budget(model_name: &str) -> usize {
116 parse_embedding_model_alias(model_name)
117 .and_then(|model| model.document_instruction())
118 .map_or(MAX_TEXT_BYTES, |prefix| {
119 MAX_TEXT_BYTES.saturating_sub(prefix.len())
120 })
121}
122
123pub fn bounded_embedding_input(text: &str, max_bytes: usize) -> (&str, bool) {
125 if text.len() <= max_bytes {
126 return (text, false);
127 }
128
129 let end = text
130 .char_indices()
131 .map(|(index, _)| index)
132 .take_while(|index| *index <= max_bytes)
133 .last()
134 .unwrap_or(0);
135 (&text[..end], true)
136}
137
138fn prepared_document_fingerprint(text: &str, model: EmbeddingModel) -> ContentRef {
139 match model.document_instruction() {
140 Some(prefix) => VectorRecord::fingerprint_text(&format!("{prefix}{text}")),
141 None => VectorRecord::fingerprint_text(text),
142 }
143}
144
145impl KhiveRuntime {
146 fn require_default_embedder(&self) -> RuntimeResult<&str> {
147 let model_name = self.default_embedder_name();
148 if model_name.is_empty() {
149 return Err(RuntimeError::Unconfigured("embedding_model".into()));
150 }
151 Ok(model_name)
152 }
153
154 pub async fn embed(&self, text: &str) -> RuntimeResult<Vec<f32>> {
159 let model_name = self.require_default_embedder()?;
160 self.embed_with_model(model_name, text).await
161 }
162
163 pub async fn embed_with_model(&self, model_name: &str, text: &str) -> RuntimeResult<Vec<f32>> {
179 let model = parse_embedding_model_alias(model_name);
180 let service = self.embedder(model_name).await?;
181 let emb_model = model.unwrap_or_default();
182 crate::usage::count(crate::usage::UsageUnit::EmbedCalls, 1);
186 let out = with_embedding_admission(service.embed_one(text, emb_model)).await;
187 out
188 }
189
190 pub async fn embed_document_with_model_outcome(
209 &self,
210 model_name: &str,
211 text: &str,
212 ) -> RuntimeResult<DocumentEmbeddingOutcome> {
213 self.embed_document_with_model_outcome_inner(None, model_name, text)
214 .await
215 }
216
217 pub(crate) async fn embed_document_with_model_outcome_for_token(
218 &self,
219 token: &NamespaceToken,
220 model_name: &str,
221 text: &str,
222 ) -> RuntimeResult<DocumentEmbeddingOutcome> {
223 self.embed_document_with_model_outcome_inner(Some(token), model_name, text)
224 .await
225 }
226
227 async fn embed_document_with_model_outcome_inner(
228 &self,
229 token: Option<&NamespaceToken>,
230 model_name: &str,
231 text: &str,
232 ) -> RuntimeResult<DocumentEmbeddingOutcome> {
233 let model = parse_embedding_model_alias(model_name);
234 let (service, audited_document_preparation) = self
235 .embedder_with_input_attestation(model_name, token)
236 .await?;
237 let emb_model = model.unwrap_or_default();
238 let source_bytes = text.len();
239 let (text, truncated) =
240 bounded_embedding_input(text, document_embedding_budget(model_name));
241 let embedded_bytes = text.len();
242 crate::usage::count(crate::usage::UsageUnit::EmbedCalls, 1);
244 let embeddings =
245 with_embedding_admission(service.embed_passage(&[text.to_string()], emb_model)).await;
246 let mut vectors = embeddings?;
247 if vectors.len() != 1 {
248 return Err(RuntimeError::Internal(format!(
249 "embed_passage returned {} vectors for 1 input",
250 vectors.len()
251 )));
252 }
253 let out = vectors.pop().expect("checked len == 1 above");
254 Ok(DocumentEmbeddingOutcome {
255 model_name: model_name.to_owned(),
256 vector: out,
257 prepared_text_fingerprint: audited_document_preparation
258 .then(|| prepared_document_fingerprint(text, emb_model)),
259 source_bytes,
260 embedded_bytes,
261 truncated,
262 })
263 }
264
265 pub async fn embed_query_with_model(
279 &self,
280 model_name: &str,
281 text: &str,
282 ) -> RuntimeResult<Vec<f32>> {
283 self.embed_query_with_model_inner(None, model_name, text)
284 .await
285 }
286
287 pub(crate) async fn embed_query_with_model_for_token(
288 &self,
289 token: &NamespaceToken,
290 model_name: &str,
291 text: &str,
292 ) -> RuntimeResult<Vec<f32>> {
293 self.embed_query_with_model_inner(Some(token), model_name, text)
294 .await
295 }
296
297 async fn embed_query_with_model_inner(
298 &self,
299 token: Option<&NamespaceToken>,
300 model_name: &str,
301 text: &str,
302 ) -> RuntimeResult<Vec<f32>> {
303 let _phase = khive_storage::request_context::enter_search_phase(
304 khive_storage::request_context::SearchPhase::VectorEmbedding,
305 );
306 let model = parse_embedding_model_alias(model_name);
307 let service = match token {
308 Some(token) => self.embedder_with_token(token, model_name).await?,
309 None => self.embedder(model_name).await?,
310 };
311 let texts = [text.to_string()];
312 let emb_model = model.unwrap_or(CUSTOM_QUERY_ROLE_MODEL);
313 crate::usage::count(crate::usage::UsageUnit::EmbedCalls, 1);
315 let embeddings = match model {
316 Some(
317 EmbeddingModel::BgeSmallEnV15
318 | EmbeddingModel::BgeBaseEnV15
319 | EmbeddingModel::BgeLargeEnV15,
320 ) => with_embedding_admission(service.embed(&texts, emb_model)).await,
321 _ => with_embedding_admission(service.embed_query(&texts, emb_model)).await,
322 };
323 let out = embeddings?
324 .into_iter()
325 .next()
326 .ok_or_else(|| RuntimeError::Internal("embed_query returned empty vec".into()))?;
327 Ok(out)
328 }
329
330 pub async fn embed_document(&self, text: &str) -> RuntimeResult<Vec<f32>> {
339 let outcome = self.embed_document_outcome(text).await?;
340 if outcome.truncated {
341 return Err(RuntimeError::InvalidInput(format!(
342 "embedding input truncated from {} to {} bytes; use embed_document_outcome to inspect the bounded vector",
343 outcome.source_bytes, outcome.embedded_bytes
344 )));
345 }
346 Ok(outcome.vector)
347 }
348
349 pub async fn embed_document_outcome(
351 &self,
352 text: &str,
353 ) -> RuntimeResult<DocumentEmbeddingOutcome> {
354 let model_name = self.require_default_embedder()?;
355 self.embed_document_with_model_outcome(model_name, text)
356 .await
357 }
358
359 pub async fn embed_query(&self, text: &str) -> RuntimeResult<Vec<f32>> {
366 let model_name = self.require_default_embedder()?;
367 self.embed_query_with_model(model_name, text).await
368 }
369
370 async fn embed_query_for_token(
371 &self,
372 token: &NamespaceToken,
373 text: &str,
374 ) -> RuntimeResult<Vec<f32>> {
375 let model_name = self.require_default_embedder()?;
376 self.embed_query_with_model_for_token(token, model_name, text)
377 .await
378 }
379
380 pub async fn embed_batch(&self, texts: &[String]) -> RuntimeResult<Vec<Vec<f32>>> {
388 if texts.is_empty() {
389 return Ok(vec![]);
390 }
391 let model_name = self.require_default_embedder()?;
392 self.embed_batch_with_model(model_name, texts).await
393 }
394
395 pub async fn embed_batch_with_model(
400 &self,
401 model_name: &str,
402 texts: &[String],
403 ) -> RuntimeResult<Vec<Vec<f32>>> {
404 if texts.is_empty() {
405 return Ok(vec![]);
406 }
407 let model = parse_embedding_model_alias(model_name);
408 let service = self.embedder(model_name).await?;
409 let emb_model = model.unwrap_or_default();
410 let out = with_embedding_admission(service.embed(texts, emb_model)).await;
411 crate::usage::count(crate::usage::UsageUnit::EmbedCalls, texts.len() as u64);
412 out
413 }
414
415 pub async fn embed_document_batch_with_model(
429 &self,
430 model_name: &str,
431 texts: &[String],
432 ) -> RuntimeResult<Vec<Vec<f32>>> {
433 if texts.is_empty() {
434 return Ok(vec![]);
435 }
436 let outcomes = self
437 .embed_document_batch_with_model_outcomes(model_name, texts)
438 .await?;
439 let mut report = EmbeddingTruncationReport::default();
440 for outcome in &outcomes {
441 report.observe(outcome);
442 }
443 if report.any_truncated() {
444 return Err(RuntimeError::InvalidInput(format!(
445 "embedding input truncated for {} documents ({} discarded bytes); use embed_document_batch_with_model_outcomes to inspect the bounded vectors",
446 report.truncated, report.discarded_bytes
447 )));
448 }
449 Ok(outcomes.into_iter().map(|outcome| outcome.vector).collect())
450 }
451
452 pub async fn embed_document_batch_with_model_outcomes(
453 &self,
454 model_name: &str,
455 texts: &[String],
456 ) -> RuntimeResult<Vec<DocumentEmbeddingOutcome>> {
457 self.embed_document_batch_with_model_outcomes_inner(None, model_name, texts)
458 .await
459 }
460
461 pub(crate) async fn embed_document_batch_with_model_outcomes_for_token(
462 &self,
463 token: &NamespaceToken,
464 model_name: &str,
465 texts: &[String],
466 ) -> RuntimeResult<Vec<DocumentEmbeddingOutcome>> {
467 self.embed_document_batch_with_model_outcomes_inner(Some(token), model_name, texts)
468 .await
469 }
470
471 async fn embed_document_batch_with_model_outcomes_inner(
472 &self,
473 token: Option<&NamespaceToken>,
474 model_name: &str,
475 texts: &[String],
476 ) -> RuntimeResult<Vec<DocumentEmbeddingOutcome>> {
477 if texts.is_empty() {
478 return Ok(vec![]);
479 }
480 let model = parse_embedding_model_alias(model_name);
481 let (service, audited_document_preparation) = self
482 .embedder_with_input_attestation(model_name, token)
483 .await?;
484 let emb_model = model.unwrap_or_default();
485 let budget = document_embedding_budget(model_name);
486 if token.is_some() {
487 crate::usage::count(crate::usage::UsageUnit::EmbedCalls, texts.len() as u64);
489 }
490 let out = if texts.iter().all(|text| text.len() <= budget) {
491 with_embedding_admission(service.embed_passage(texts, emb_model)).await
492 } else {
493 let bounded_texts: Vec<String> = texts
494 .iter()
495 .map(|text| bounded_embedding_input(text, budget).0.to_owned())
496 .collect();
497 with_embedding_admission(service.embed_passage(&bounded_texts, emb_model)).await
498 };
499 if token.is_none() {
500 crate::usage::count(crate::usage::UsageUnit::EmbedCalls, texts.len() as u64);
501 }
502 let vectors = out?;
503 if vectors.len() != texts.len() {
504 return Err(RuntimeError::Internal(format!(
505 "embed_passage returned {} vectors for {} inputs",
506 vectors.len(),
507 texts.len()
508 )));
509 }
510 Ok(texts
511 .iter()
512 .zip(vectors)
513 .map(|(text, vector)| {
514 let (bounded, truncated) = bounded_embedding_input(text, budget);
515 DocumentEmbeddingOutcome {
516 model_name: model_name.to_owned(),
517 vector,
518 prepared_text_fingerprint: audited_document_preparation
519 .then(|| prepared_document_fingerprint(bounded, emb_model)),
520 source_bytes: text.len(),
521 embedded_bytes: bounded.len(),
522 truncated,
523 }
524 })
525 .collect())
526 }
527
528 pub async fn embed_document_batch(&self, texts: &[String]) -> RuntimeResult<Vec<Vec<f32>>> {
537 if texts.is_empty() {
538 return Ok(vec![]);
539 }
540 let model_name = self.require_default_embedder()?;
541 self.embed_document_batch_with_model(model_name, texts)
542 .await
543 }
544
545 pub async fn embed_document_batch_outcomes(
547 &self,
548 texts: &[String],
549 ) -> RuntimeResult<Vec<DocumentEmbeddingOutcome>> {
550 if texts.is_empty() {
551 return Ok(vec![]);
552 }
553 let model_name = self.require_default_embedder()?;
554 self.embed_document_batch_with_model_outcomes(model_name, texts)
555 .await
556 }
557
558 pub async fn embed_query_batch_with_model(
565 &self,
566 model_name: &str,
567 texts: &[String],
568 ) -> RuntimeResult<Vec<Vec<f32>>> {
569 if texts.is_empty() {
570 return Ok(vec![]);
571 }
572 let model = parse_embedding_model_alias(model_name);
573 let service = self.embedder(model_name).await?;
574 let emb_model = model.unwrap_or(CUSTOM_QUERY_ROLE_MODEL);
575 let out = match model {
576 Some(
577 EmbeddingModel::BgeSmallEnV15
578 | EmbeddingModel::BgeBaseEnV15
579 | EmbeddingModel::BgeLargeEnV15,
580 ) => with_embedding_admission(service.embed(texts, emb_model)).await,
581 _ => with_embedding_admission(service.embed_query(texts, emb_model)).await,
582 };
583 crate::usage::count(crate::usage::UsageUnit::EmbedCalls, texts.len() as u64);
584 out
585 }
586
587 pub(crate) async fn note_search_vector_search(
591 &self,
592 token: &NamespaceToken,
593 query_embedding: Option<Vec<f32>>,
594 query_text: &str,
595 top_k: u32,
596 ) -> RuntimeResult<Vec<VectorSearchHit>> {
597 let embedding = match query_embedding {
598 Some(embedding) => embedding,
599 None => self.embed_query_for_token(token, query_text).await?,
600 };
601 let model = self.default_embedder_name();
602 if !model.is_empty() {
603 if let Some(provider) = self.note_search_ann_provider()? {
604 if let Some(hits) = provider.search(token, model, &embedding, top_k).await? {
605 crate::note_search_ann::record_ann_route();
606 crate::usage::count(crate::usage::UsageUnit::VectorPasses, 1);
607 return Ok(hits);
608 }
609 }
610 }
611 crate::note_search_ann::record_fallback_route();
612 self.vector_search(
613 token,
614 Some(embedding),
615 None,
616 top_k,
617 Some(SubstrateKind::Note),
618 )
619 .await
620 }
621
622 #[allow(clippy::too_many_arguments)]
668 pub async fn hybrid_search(
669 &self,
670 token: &NamespaceToken,
671 query_text: &str,
672 query_vector: Option<Vec<f32>>,
673 limit: u32,
674 entity_kind: Option<&str>,
675 entity_type: Option<&str>,
676 tags_any: &[String],
677 properties_filter: Option<&serde_json::Value>,
678 ) -> RuntimeResult<Vec<SearchHit>> {
679 self.hybrid_search_with_text_mode(
680 token,
681 query_text,
682 query_vector,
683 limit,
684 entity_kind,
685 entity_type,
686 tags_any,
687 properties_filter,
688 TextQueryMode::Plain,
689 )
690 .await
691 }
692
693 #[allow(clippy::too_many_arguments)]
695 pub async fn hybrid_search_with_text_mode(
696 &self,
697 token: &NamespaceToken,
698 query_text: &str,
699 query_vector: Option<Vec<f32>>,
700 limit: u32,
701 entity_kind: Option<&str>,
702 entity_type: Option<&str>,
703 tags_any: &[String],
704 properties_filter: Option<&serde_json::Value>,
705 text_mode: TextQueryMode,
706 ) -> RuntimeResult<Vec<SearchHit>> {
707 let (hits, _vector_error) = self
708 .hybrid_search_inner(
709 token,
710 query_text,
711 query_vector,
712 limit,
713 entity_kind,
714 entity_type,
715 tags_any,
716 properties_filter,
717 text_mode,
718 None,
719 false,
720 None,
721 )
722 .await?;
723 Ok(hits)
724 }
725
726 pub async fn hybrid_search_each_kind(
737 &self,
738 token: &NamespaceToken,
739 query_text: &str,
740 query_vector: Option<Vec<f32>>,
741 limit: u32,
742 entity_kinds: &[&str],
743 ) -> RuntimeResult<Vec<Vec<SearchHit>>> {
744 if entity_kinds.is_empty() {
745 return Ok(Vec::new());
746 }
747 let candidates = limit.saturating_mul(CANDIDATE_MULTIPLIER).max(limit);
748 let (vector_hits, _vector_error) = self
749 .hybrid_vector_stage(token, query_text, query_vector, candidates, None, false)
750 .await?;
751 let mut per_kind = Vec::with_capacity(entity_kinds.len());
752 for &kind in entity_kinds {
753 let (hits, _vector_error) = self
754 .hybrid_search_inner(
755 token,
756 query_text,
757 None,
758 limit,
759 Some(kind),
760 None,
761 &[],
762 None,
763 TextQueryMode::Plain,
764 None,
765 false,
766 Some(vector_hits.clone()),
767 )
768 .await?;
769 per_kind.push(hits);
770 }
771 Ok(per_kind)
772 }
773
774 #[allow(clippy::too_many_arguments)]
779 pub(crate) async fn hybrid_search_with_vector_similarity_floor(
780 &self,
781 token: &NamespaceToken,
782 query_text: &str,
783 query_vector: Option<Vec<f32>>,
784 limit: u32,
785 entity_kind: Option<&str>,
786 entity_type: Option<&str>,
787 tags_any: &[String],
788 properties_filter: Option<&serde_json::Value>,
789 vector_similarity_floor: f64,
790 ) -> RuntimeResult<Vec<SearchHit>> {
791 let (hits, _vector_error) = self
792 .hybrid_search_inner(
793 token,
794 query_text,
795 query_vector,
796 limit,
797 entity_kind,
798 entity_type,
799 tags_any,
800 properties_filter,
801 TextQueryMode::Plain,
802 Some(vector_similarity_floor),
803 false,
804 None,
805 )
806 .await?;
807 Ok(hits)
808 }
809
810 #[allow(clippy::too_many_arguments)]
821 pub async fn hybrid_search_outcome(
822 &self,
823 token: &NamespaceToken,
824 query_text: &str,
825 limit: u32,
826 entity_kind: Option<&str>,
827 entity_type: Option<&str>,
828 tags_any: &[String],
829 properties_filter: Option<&serde_json::Value>,
830 ) -> RuntimeResult<HybridSearchOutcome> {
831 self.hybrid_search_outcome_with_text_mode(
832 token,
833 query_text,
834 limit,
835 entity_kind,
836 entity_type,
837 tags_any,
838 properties_filter,
839 TextQueryMode::Plain,
840 )
841 .await
842 }
843
844 #[allow(clippy::too_many_arguments)]
846 pub async fn hybrid_search_outcome_with_text_mode(
847 &self,
848 token: &NamespaceToken,
849 query_text: &str,
850 limit: u32,
851 entity_kind: Option<&str>,
852 entity_type: Option<&str>,
853 tags_any: &[String],
854 properties_filter: Option<&serde_json::Value>,
855 text_mode: TextQueryMode,
856 ) -> RuntimeResult<HybridSearchOutcome> {
857 let (hits, vector_error) = self
858 .hybrid_search_inner(
859 token,
860 query_text,
861 None,
862 limit,
863 entity_kind,
864 entity_type,
865 tags_any,
866 properties_filter,
867 text_mode,
868 None,
869 true,
870 None,
871 )
872 .await?;
873 Ok(HybridSearchOutcome { hits, vector_error })
874 }
875
876 #[allow(clippy::too_many_arguments)]
879 async fn hybrid_search_inner(
880 &self,
881 token: &NamespaceToken,
882 query_text: &str,
883 query_vector: Option<Vec<f32>>,
884 limit: u32,
885 entity_kind: Option<&str>,
886 entity_type: Option<&str>,
887 tags_any: &[String],
888 properties_filter: Option<&serde_json::Value>,
889 text_mode: TextQueryMode,
890 vector_similarity_floor: Option<f64>,
891 tolerate_vector_error: bool,
892 vector_pool: Option<Vec<VectorSearchHit>>,
893 ) -> RuntimeResult<(Vec<SearchHit>, Option<String>)> {
894 khive_storage::await_request_read_phase("hybrid_search", async {
895 let candidates = limit.saturating_mul(CANDIDATE_MULTIPLIER).max(limit);
896
897 let visible_ns: Vec<String> = token
898 .visible_namespaces()
899 .iter()
900 .map(|ns| ns.as_str().to_owned())
901 .collect();
902 let text_store = self.text(token)?;
906 let text_fut = async {
907 let _phase = khive_storage::request_context::enter_search_phase(
908 khive_storage::request_context::SearchPhase::TextSearch,
909 );
910 text_store
911 .search(TextSearchRequest {
912 query: query_text.to_string(),
913 mode: text_mode,
914 filter: Some(TextFilter {
915 namespaces: visible_ns.clone(),
916 record_kinds: entity_kind
923 .map(|kind| vec![kind.to_string()])
924 .unwrap_or_default(),
925 ..TextFilter::default()
926 }),
927 top_k: candidates,
928 snippet_chars: 200,
929 })
930 .await
931 };
932 let text_fut = crate::stage_seam::text_stage(text_fut);
933 let vector_fut = async {
935 let _phase = khive_storage::request_context::enter_search_phase(
936 khive_storage::request_context::SearchPhase::VectorSearch,
937 );
938 match vector_pool {
939 Some(pool) => Ok((pool, None)),
940 None => {
941 self.hybrid_vector_stage(
942 token,
943 query_text,
944 query_vector,
945 candidates,
946 vector_similarity_floor,
947 tolerate_vector_error,
948 )
949 .await
950 }
951 }
952 };
953 let (text_search_result, vector_result) = tokio::join!(text_fut, vector_fut);
954 let text_hits = crate::error::fts_text_leg_or_err(
959 text_search_result.map_err(RuntimeError::from),
960 "hybrid_search",
961 query_text,
962 )?;
963 let (vector_hits, vector_error) = vector_result?;
964
965 let fusion_limit = text_hits.len().saturating_add(vector_hits.len());
970 let mut fused = {
971 let _phase = khive_storage::request_context::enter_search_phase(
972 khive_storage::request_context::SearchPhase::Fusion,
973 );
974 rrf_fuse(text_hits, vector_hits, fusion_limit, query_text)
975 };
976
977 if !fused.is_empty() {
980 let candidate_ids: Vec<Uuid> = fused.iter().map(|h| h.entity_id).collect();
981 let _phase = khive_storage::request_context::enter_search_phase(
982 khive_storage::request_context::SearchPhase::EntityHydration,
983 );
984 let alive_page = self
985 .entities(token)?
986 .query_entities(
987 token.namespace().as_str(),
988 EntityFilter {
989 ids: candidate_ids,
990 kinds: entity_kind.map(|k| vec![k.to_string()]).unwrap_or_default(),
991 entity_types: entity_type
992 .map(|t| vec![t.to_string()])
993 .unwrap_or_default(),
994 namespaces: visible_ns,
995 tags_any: tags_any.to_vec(),
996 ..EntityFilter::default()
997 },
998 PageRequest {
999 offset: 0,
1000 limit: u32::try_from(fused.len()).unwrap_or(u32::MAX),
1001 },
1002 )
1003 .await?;
1004 drop(_phase);
1005 let mut entity_meta: HashMap<Uuid, (String, Option<String>)> = HashMap::new();
1006 let mut alive: HashSet<Uuid> = HashSet::new();
1007 for e in alive_page.items {
1008 if let Some(pf) = properties_filter {
1011 if !properties_match(e.properties.as_ref(), pf) {
1012 continue;
1013 }
1014 }
1015 alive.insert(e.id);
1016 entity_meta.insert(e.id, (e.name, e.description));
1017 }
1018
1019 fused.retain(|h| alive.contains(&h.entity_id));
1020
1021 for hit in &mut fused {
1023 if let Some((name, description)) = entity_meta.get(&hit.entity_id) {
1024 if hit.title.is_none() {
1025 hit.title = Some(name.clone());
1026 }
1027 if hit.snippet.is_none() {
1028 hit.snippet = description.clone();
1029 }
1030 }
1031 }
1032 }
1033
1034 fused.truncate(limit as usize);
1035 Ok((fused, vector_error))
1036 })
1037 .await?
1038 }
1039
1040 async fn embed_backfill_page(
1041 &self,
1042 token: &NamespaceToken,
1043 model_name: &str,
1044 inputs: &[(Uuid, String)],
1045 ) -> Vec<Option<DocumentEmbeddingOutcome>> {
1046 let texts: Vec<String> = inputs.iter().map(|(_, text)| text.clone()).collect();
1047 match self
1048 .embed_document_batch_with_model_outcomes_for_token(token, model_name, &texts)
1049 .await
1050 {
1051 Ok(outcomes) => outcomes.into_iter().map(Some).collect(),
1052 Err(error) => {
1053 tracing::warn!(
1054 model = %model_name,
1055 error = %error,
1056 "backfill_missing_embeddings: batch embed failed; retrying records individually"
1057 );
1058 let mut outcomes = Vec::with_capacity(inputs.len());
1059 for (id, text) in inputs {
1060 match self
1061 .embed_document_with_model_outcome_for_token(token, model_name, text)
1062 .await
1063 {
1064 Ok(outcome) => outcomes.push(Some(outcome)),
1065 Err(error) => {
1066 tracing::warn!(
1067 id = %id,
1068 model = %model_name,
1069 error = %error,
1070 "backfill_missing_embeddings: record embed failed"
1071 );
1072 outcomes.push(None);
1073 }
1074 }
1075 }
1076 outcomes
1077 }
1078 }
1079 }
1080
1081 pub async fn backfill_missing_embeddings(&self, token: &NamespaceToken) -> RuntimeResult<u64> {
1094 use khive_storage::types::{SqlRow, SqlStatement, SqlValue};
1095
1096 let model_names = self.registered_embedding_model_names();
1097 if model_names.is_empty() {
1098 tracing::debug!(
1099 "backfill_missing_embeddings: no embedding models registered, skipping"
1100 );
1101 return Ok(0);
1102 }
1103
1104 let ns = token.namespace().as_str().to_string();
1105 let mut total_backfilled = 0u64;
1106
1107 for model_name in &model_names {
1108 let mut model_truncation = EmbeddingTruncationReport::default();
1109 match self.vectors_for_model(token, model_name) {
1110 Ok(_) => {}
1111 Err(error) => {
1112 tracing::warn!(model = %model_name, error = %error,
1113 "backfill_missing_embeddings: vector store unavailable");
1114 continue;
1115 }
1116 };
1117 let vec_table = format!("vec_{}", sanitize_key(model_name));
1119
1120 const PAGE_SIZE: usize = EMBEDDING_BATCH_PAGE_SIZE;
1124 let mut entity_total = 0usize;
1125 let mut entity_cursor = String::new();
1126 loop {
1127 let entity_sql = SqlStatement {
1128 sql: format!(
1129 "SELECT id FROM entities \
1130 WHERE namespace = ?1 AND deleted_at IS NULL AND id > ?3 \
1131 AND id NOT IN (\
1132 SELECT subject_id FROM {vec_table} \
1133 WHERE namespace = ?1 AND embedding_model = ?2 \
1134 ) ORDER BY id LIMIT {PAGE_SIZE}"
1135 ),
1136 params: vec![
1137 SqlValue::Text(ns.clone()),
1138 SqlValue::Text(model_name.clone()),
1139 SqlValue::Text(entity_cursor.clone()),
1140 ],
1141 label: Some("backfill_entities".into()),
1142 };
1143
1144 let entity_rows: Vec<SqlRow> = {
1145 let sql = self.sql();
1146 let reader_result = sql.reader().await;
1147 #[cfg(any(test, feature = "fault-injection"))]
1148 let reader_result = if BACKFILL_READER_FAIL.with(|c| c.get()) {
1149 BACKFILL_READER_FAIL.with(|c| c.set(false));
1150 Err(khive_storage::StorageError::Pool {
1151 operation: "reader".into(),
1152 message: "injected failure".into(),
1153 })
1154 } else {
1155 reader_result
1156 };
1157 let mut reader = reader_result.map_err(RuntimeError::Storage)?;
1158 reader
1159 .query_all(entity_sql)
1160 .await
1161 .map_err(RuntimeError::Storage)?
1162 };
1163
1164 let batch_len = entity_rows.len();
1165 entity_total += batch_len;
1166 if batch_len == 0 {
1167 break;
1168 }
1169 entity_cursor = match entity_rows.last().and_then(|row| row.columns.first()) {
1170 Some(column) => match &column.value {
1171 SqlValue::Text(id) => id.clone(),
1172 _ => {
1173 return Err(RuntimeError::Internal(
1174 "backfill entity ID is not text".into(),
1175 ))
1176 }
1177 },
1178 None => {
1179 return Err(RuntimeError::Internal(
1180 "backfill entity page is empty".into(),
1181 ))
1182 }
1183 };
1184
1185 let entity_store = self.entities(token)?;
1186 let mut entities = Vec::with_capacity(batch_len);
1187 let mut inputs = Vec::with_capacity(batch_len);
1188 for row in &entity_rows {
1189 let Some(SqlValue::Text(id)) = row.columns.first().map(|column| &column.value)
1190 else {
1191 continue;
1192 };
1193 let Ok(id) = id.parse::<Uuid>() else { continue };
1194 let Some(entity) = entity_store.get_entity(id).await? else {
1195 continue;
1196 };
1197 if entity.namespace != ns || entity.deleted_at.is_some() {
1198 continue;
1199 }
1200 let text = crate::curation::entity_embedding_text(&entity);
1201 if text.trim().is_empty() {
1202 continue;
1203 }
1204 inputs.push((id, text));
1205 entities.push(entity);
1206 }
1207 let outcomes = self.embed_backfill_page(token, model_name, &inputs).await;
1208 for (entity, outcome) in entities.into_iter().zip(outcomes) {
1209 let Some(outcome) = outcome else { continue };
1210 model_truncation.observe(&outcome);
1211 match self
1212 .publish_entity_vector_revision(token, &entity, model_name, &outcome.vector)
1213 .await
1214 {
1215 Ok(true) => total_backfilled += 1,
1216 Ok(false) => {}
1217 Err(error) => tracing::warn!(
1218 id = %entity.id, model = %model_name, error = %error,
1219 "backfill_missing_embeddings: entity vector insert failed"
1220 ),
1221 }
1222 }
1223
1224 if batch_len < PAGE_SIZE {
1225 break;
1226 }
1227 }
1228
1229 let text_store = self.text_for_notes(token).ok();
1231 let note_store = self.notes(token).ok();
1232 let mut note_total = 0usize;
1233 let mut note_cursor = String::new();
1236 loop {
1237 let note_sql = SqlStatement {
1240 sql: format!(
1241 "SELECT id FROM notes \
1242 WHERE namespace = ?1 AND deleted_at IS NULL AND id > ?3 \
1243 AND id NOT IN (\
1244 SELECT subject_id FROM {vec_table} \
1245 WHERE namespace = ?1 AND embedding_model = ?2 \
1246 ) ORDER BY id LIMIT {PAGE_SIZE}"
1247 ),
1248 params: vec![
1249 SqlValue::Text(ns.clone()),
1250 SqlValue::Text(model_name.clone()),
1251 SqlValue::Text(note_cursor.clone()),
1252 ],
1253 label: Some("backfill_notes".into()),
1254 };
1255
1256 let note_rows: Vec<SqlRow> = {
1257 let sql = self.sql();
1258 let reader_result = sql.reader().await;
1259 #[cfg(any(test, feature = "fault-injection"))]
1260 let reader_result = if BACKFILL_READER_FAIL.with(|c| c.get()) {
1261 BACKFILL_READER_FAIL.with(|c| c.set(false));
1262 Err(khive_storage::StorageError::Pool {
1263 operation: "reader".into(),
1264 message: "injected failure".into(),
1265 })
1266 } else {
1267 reader_result
1268 };
1269 let mut reader = reader_result.map_err(RuntimeError::Storage)?;
1270 reader
1271 .query_all(note_sql)
1272 .await
1273 .map_err(RuntimeError::Storage)?
1274 };
1275
1276 let batch_len = note_rows.len();
1277 note_total += batch_len;
1278 if batch_len == 0 {
1279 break;
1280 }
1281 note_cursor = match note_rows.last().and_then(|row| row.columns.first()) {
1282 Some(column) => match &column.value {
1283 SqlValue::Text(id) => id.clone(),
1284 _ => {
1285 return Err(RuntimeError::Internal(
1286 "backfill note ID is not text".into(),
1287 ))
1288 }
1289 },
1290 None => {
1291 return Err(RuntimeError::Internal("backfill note page is empty".into()))
1292 }
1293 };
1294
1295 let mut notes_for_batch = Vec::with_capacity(batch_len);
1296 let mut inputs = Vec::with_capacity(batch_len);
1297 for row in ¬e_rows {
1298 let id_str = row.columns.first().and_then(|c| {
1299 if let SqlValue::Text(s) = &c.value {
1300 Some(s.clone())
1301 } else {
1302 None
1303 }
1304 });
1305
1306 let Some(id_str) = id_str else {
1307 continue;
1308 };
1309 let Ok(id) = id_str.parse::<Uuid>() else {
1310 continue;
1311 };
1312
1313 let note = match ¬e_store {
1314 Some(store) => match store.get_note(id).await {
1315 Ok(Some(n)) => n,
1316 _ => continue,
1317 },
1318 None => continue,
1319 };
1320
1321 if note.content.trim().is_empty() {
1322 continue;
1323 }
1324
1325 if model_names.first().map(|n| n.as_str()) == Some(model_name.as_str()) {
1328 if let Some(ref ts) = text_store {
1329 if let Err(e) = ts.upsert_document(note_fts_document(¬e)).await {
1330 tracing::warn!(id = %id, error = %e,
1331 "backfill_missing_embeddings: note FTS upsert failed");
1332 }
1333 }
1334 }
1335
1336 if !self
1337 .embedding_models_for_note_kind(¬e.kind)
1338 .contains(model_name)
1339 {
1340 continue;
1341 }
1342
1343 inputs.push((id, note.content.clone()));
1344 notes_for_batch.push(note);
1345 }
1346
1347 let outcomes = self.embed_backfill_page(token, model_name, &inputs).await;
1348 for (note, outcome) in notes_for_batch.into_iter().zip(outcomes) {
1349 let Some(outcome) = outcome else { continue };
1350 model_truncation.observe(&outcome);
1351 match self
1352 .publish_note_vector_revision(token, ¬e, model_name, &outcome.vector)
1353 .await
1354 {
1355 Ok(true) => total_backfilled += 1,
1356 Ok(false) => {}
1357 Err(error) => tracing::warn!(
1358 id = %note.id, model = %model_name, error = %error,
1359 "backfill_missing_embeddings: note vector insert failed"
1360 ),
1361 }
1362 }
1363
1364 if batch_len < PAGE_SIZE {
1365 break;
1366 }
1367 }
1368
1369 tracing::info!(
1370 model = %model_name,
1371 namespace = %ns,
1372 entities = entity_total,
1373 notes = note_total,
1374 truncated = model_truncation.truncated,
1375 discarded_bytes = model_truncation.discarded_bytes,
1376 "backfill_missing_embeddings: model pass complete"
1377 );
1378 }
1379
1380 tracing::info!(
1381 namespace = %ns,
1382 total_backfilled = total_backfilled,
1383 "backfill_missing_embeddings: finished"
1384 );
1385
1386 Ok(total_backfilled)
1387 }
1388
1389 pub async fn sweep_orphan_vectors(
1405 &self,
1406 token: &NamespaceToken,
1407 max_delete_per_model: u32,
1408 dry_run: bool,
1409 ) -> RuntimeResult<u64> {
1410 use khive_storage::types::OrphanSweepConfig;
1411 use khive_storage::StorageError;
1412
1413 let model_names = self.registered_embedding_model_names();
1414 if model_names.is_empty() {
1415 tracing::debug!("sweep_orphan_vectors: no embedding models registered, skipping");
1416 return Ok(0);
1417 }
1418
1419 let ns = token.namespace().as_str().to_string();
1420 let mut total_deleted = 0u64;
1421
1422 for model_name in &model_names {
1423 let store = match self.vectors_for_model(token, model_name) {
1424 Ok(s) => s,
1425 Err(e) => {
1426 tracing::warn!(
1427 model = %model_name,
1428 error = %e,
1429 "sweep_orphan_vectors: failed to get vector store, skipping model"
1430 );
1431 continue;
1432 }
1433 };
1434
1435 let caps = store.capabilities();
1436 if !caps.supports_orphan_sweep {
1437 tracing::debug!(
1438 model = %model_name,
1439 "sweep_orphan_vectors: backend does not support orphan sweep, skipping"
1440 );
1441 continue;
1442 }
1443
1444 let config = OrphanSweepConfig {
1445 subject_id_allowlist: None,
1446 namespaces: vec![ns.clone()],
1447 substrate_kinds: vec![],
1448 max_delete: max_delete_per_model,
1449 dry_run,
1450 };
1451
1452 match store.orphan_sweep(&config).await {
1453 Ok(result) => {
1454 tracing::info!(
1455 model = %model_name,
1456 namespace = %ns,
1457 scanned = result.scanned,
1458 deleted = result.deleted,
1459 would_delete = result.would_delete,
1460 dry_run = dry_run,
1461 "sweep_orphan_vectors: sweep complete"
1462 );
1463 total_deleted += result.deleted;
1464 }
1465 Err(StorageError::Unsupported { .. }) => {
1466 tracing::debug!(
1467 model = %model_name,
1468 "sweep_orphan_vectors: backend returned Unsupported, skipping"
1469 );
1470 }
1471 Err(e) => {
1472 tracing::warn!(
1473 model = %model_name,
1474 error = %e,
1475 "sweep_orphan_vectors: sweep failed, continuing with other models"
1476 );
1477 }
1478 }
1479 }
1480
1481 tracing::info!(
1482 namespace = %ns,
1483 total_deleted = total_deleted,
1484 dry_run = dry_run,
1485 "sweep_orphan_vectors: finished"
1486 );
1487
1488 Ok(total_deleted)
1489 }
1490}
1491
1492pub(crate) fn properties_match(
1497 properties: Option<&serde_json::Value>,
1498 filter: &serde_json::Value,
1499) -> bool {
1500 let required = match filter.as_object() {
1501 Some(obj) if !obj.is_empty() => obj,
1502 _ => return true,
1503 };
1504 let actual = match properties.and_then(serde_json::Value::as_object) {
1505 Some(obj) => obj,
1506 None => return false,
1507 };
1508 required
1509 .iter()
1510 .all(|(k, v)| actual.get(k).is_some_and(|av| av == v))
1511}
1512
1513const EXACT_MATCH_BOOST: f64 = 0.5;
1517
1518fn rrf_fuse(
1527 text_hits: Vec<TextSearchHit>,
1528 vector_hits: Vec<VectorSearchHit>,
1529 limit: usize,
1530 query_text: &str,
1531) -> Vec<SearchHit> {
1532 let mut text_arm = Vec::new();
1533 for (rank, hit) in text_hits.into_iter().enumerate() {
1534 let label = HitLabel {
1535 rank,
1536 signals: SearchSignals {
1537 vector_similarity: None,
1538 keyword_score: Some(hit.score),
1539 },
1540 source: SearchSource::Text,
1541 title: hit.title,
1542 snippet: hit.snippet,
1543 };
1544 text_arm.push((hit.subject_id, label));
1545 }
1546 let mut vector_arm = Vec::new();
1547 for (rank, hit) in vector_hits.into_iter().enumerate() {
1548 let label = HitLabel {
1549 rank,
1550 signals: SearchSignals {
1551 vector_similarity: Some(hit.score),
1552 keyword_score: None,
1553 },
1554 source: SearchSource::Vector,
1555 title: None,
1556 snippet: None,
1557 };
1558 vector_arm.push((hit.subject_id, label));
1559 }
1560
1561 let fused = fuse_labelled(
1562 vec![text_arm, vector_arm],
1563 RRF_K,
1564 combine_leg_first_appearance,
1565 );
1566
1567 let query_lower = query_text.to_lowercase();
1569 let boost = DeterministicScore::from_f64(EXACT_MATCH_BOOST);
1570 let mut hits = Vec::with_capacity(fused.len());
1571 for (entity_id, score, label) in fused {
1572 let title_lower = label.title.as_deref().map(str::to_lowercase);
1573 let exact = title_lower.as_deref() == Some(query_lower.as_str());
1574 let score = if exact { score + boost } else { score };
1575 hits.push(SearchHit {
1576 entity_id,
1577 score,
1578 rank_score_kind: RankScoreKind::Rrf,
1579 signals: label.signals,
1580 source: label.source,
1581 title: label.title,
1582 snippet: label.snippet,
1583 });
1584 }
1585
1586 hits.sort_by(|a, b| {
1587 khive_score::cmp_desc_then_id(a.score, &a.entity_id, b.score, &b.entity_id)
1588 });
1589 hits.truncate(limit);
1590 hits
1591}
1592
1593#[cfg(test)]
1594mod rrf_fuse_label_tests;
1595
1596#[cfg(test)]
1597#[path = "retrieval_tests.rs"]
1598mod tests;