Skip to main content

khive_runtime/
retrieval.rs

1//! Retrieval operations: local embedding generation and hybrid search with RRF fusion.
2
3use 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
23/// The model value a custom provider's query-role method receives. A custom
24/// provider has no built-in model, but the service trait takes one, and the
25/// trait's default query method prepends that model's query instruction. This
26/// model has neither a query nor a document instruction, so a custom query
27/// keeps its original text whether or not the provider overrides the method.
28const 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
39/// Bounds provider input and per-page outcome memory while amortizing model setup.
40pub(crate) const EMBEDDING_BATCH_PAGE_SIZE: usize = 256;
41
42// Fault-injection flag for backfill reader errors (test / `fault-injection` builds only).
43#[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/// Arm the backfill reader fault injection: the next `backfill_missing_embeddings`
49/// call substitutes a `StorageError::Pool` for `sql.reader().await`'s result, then
50/// resets the flag. The injected error passes through the same
51/// `map_err(RuntimeError::Storage)?` path as a real reader failure, exercising the
52/// fail-closed guard rather than bypassing it.
53#[cfg(any(test, feature = "fault-injection"))]
54pub fn arm_backfill_reader_fail() {
55    BACKFILL_READER_FAIL.with(|c| c.set(true));
56}
57
58/// RRF constant. Controls how strongly top ranks dominate.
59///
60/// The paper's k=60 over-compresses scores at KG scale (tens–thousands of
61/// entities): rank 1 ≈ 0.016, rank 10 ≈ 0.014, spread ≈ 0.002. k=10 gives
62/// rank 1 ≈ 0.091, rank 10 ≈ 0.050, spread ≈ 0.041 — 20× better discrimination,
63/// which dedup-before-create needs at graph sizes of 50–2700 entities.
64const RRF_K: usize = 10;
65
66/// Candidates pulled per path before fusion. Higher = better recall, more work.
67const CANDIDATE_MULTIPLIER: u32 = 4;
68
69/// Advisory emitted by write verbs when only the embedding input was bounded.
70pub const EMBEDDING_INPUT_TRUNCATED_WARNING: &str =
71    "embedding input was truncated to the embedder maximum; full content was stored unchanged";
72
73/// Outcome for one document embedding. The exact-input fingerprint is present
74/// only when the resolved provider has the runtime-audited preparation path.
75#[derive(Clone, Debug)]
76pub struct DocumentEmbeddingOutcome {
77    pub model_name: String,
78    pub vector: Vec<f32>,
79    /// Digest of the exact prepared input, or `None` for an unaudited provider.
80    pub prepared_text_fingerprint: Option<ContentRef>,
81    pub source_bytes: usize,
82    pub embedded_bytes: usize,
83    pub truncated: bool,
84}
85
86/// Aggregate truncation accounting for one logical write or reindex batch.
87#[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    /// Add another batch's counters without wrapping on overflow.
108    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
114/// Maximum document bytes accepted before the embedding service adds its model prefix.
115pub 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
123/// Return the longest UTF-8-safe prefix of `text` within `max_bytes` and whether it was shortened.
124pub 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    /// Generate an embedding vector for `text` using the configured default model.
155    ///
156    /// First call lazily loads model weights (cold start cost). Subsequent calls reuse them.
157    /// Returns `Unconfigured("embedding_model")` if no model is configured.
158    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    /// Generate an embedding vector for `text` using the named model.
164    ///
165    /// Accepts both built-in lattice model names/aliases and custom provider
166    /// names registered via [`KhiveRuntime::register_embedder`]. For lattice
167    /// models the resolved `EmbeddingModel` enum is forwarded to `embed_one`
168    /// so the service can select the correct model variant. For custom
169    /// providers, `embed_one` is called with `EmbeddingModel::default()`
170    /// because custom services are expected to ignore the enum argument (they
171    /// own a single model implicitly).
172    ///
173    /// Applies no instruction prefix (generic role). Use
174    /// [`Self::embed_document_with_model_outcome`] / [`Self::embed_query_with_model`] for
175    /// instruction-tuned models where the asymmetric prefix matters.
176    ///
177    /// Returns `UnknownModel` if `model_name` is not in the embedder registry.
178    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        // Issued-at-dispatch: count before the provider await so a call that
183        // was handed to the provider is counted even if this task is aborted
184        // while parked on the await (drain_embed_join_set cancellation path).
185        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    /// Embed a document/passage for indexing using the named model.
191    ///
192    /// Applies `EmbeddingService::embed_passage`, which prepends the model's
193    /// `document_instruction()` prefix when defined (e.g. `"passage: "` for
194    /// multilingual-e5). For models with no document prefix (MiniLM, BGE) this
195    /// is identical to [`Self::embed_with_model`].
196    ///
197    /// Use this for all index/store/backfill paths so that instruction-tuned
198    /// models produce passage-side vectors.
199    ///
200    /// **Reindex caveat**: switching from an unprefixed model (or a model with no
201    /// `document_instruction`) to an instruction-tuned model changes the vector
202    /// representation. Vectors stored under the old scheme are not comparable to
203    /// newly prefixed vectors. Operators must trigger a full reindex
204    /// (`knowledge.index(rebuild_ann=true)` / `kkernel reindex`) after changing
205    /// the embedding model config.
206    ///
207    /// Returns `UnknownModel` if `model_name` is not registered.
208    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        // Issued-at-dispatch: counted before the await — see embed_with_model.
243        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    /// Embed a query string for retrieval using the named model.
266    ///
267    /// Applies the pre-0.6 query-role behavior. E5 and Qwen models retain
268    /// their query instructions and BGE receives the original unprefixed query
269    /// text. A custom provider receives the original text through its
270    /// query-role method and owns its own query preparation; the trait's
271    /// default query method adds no instruction to it.
272    ///
273    /// Use this for all search/recall/suggest query embedding paths so that
274    /// instruction-tuned models land in the correct side of their retrieval
275    /// space.
276    ///
277    /// Returns `UnknownModel` if `model_name` is not registered.
278    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        // Issued-at-dispatch: counted before the await — see embed_with_model.
314        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    /// Embed a document for indexing using the configured default model.
331    ///
332    /// Delegates to [`Self::embed_document_outcome`]. Use for entity/note
333    /// create and reindex paths.
334    ///
335    /// Returns `Unconfigured("embedding_model")` if no model is configured.
336    /// Returns an error if the input is bounded; use the outcome method to
337    /// inspect the vector together with the truncation metadata.
338    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    /// Embed a document with the default model and retain input-bounding metadata.
350    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    /// Embed a query for retrieval using the configured default model.
360    ///
361    /// Delegates to [`Self::embed_query_with_model`]. Use for vector search and
362    /// hybrid search query paths.
363    ///
364    /// Returns `Unconfigured("embedding_model")` if no model is configured.
365    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    /// Generate embeddings for multiple texts in one call using the configured default model.
381    ///
382    /// Delegates to the cached `EmbeddingService::embed`, so repeated texts within
383    /// and across calls benefit from the runtime-level LRU cache.
384    ///
385    /// Returns an empty vec for empty input without hitting the embedding service.
386    /// Returns `Unconfigured("embedding_model")` if no model is configured.
387    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    /// Generate embeddings for multiple texts using the named model.
396    ///
397    /// Accepts lattice model names/aliases and custom provider names.
398    /// Returns `UnknownModel` if `model_name` is not in the embedder registry.
399    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    /// Embed a batch of documents for indexing using the named model.
416    ///
417    /// Applies `EmbeddingService::embed_passage`. Use for all bulk
418    /// index/backfill/reindex operations to apply the passage-side prefix.
419    /// A mixed batch is bounded into one ordered owned batch so truncation never
420    /// fragments one provider batch into sequential singleton inference calls.
421    ///
422    /// **Reindex caveat**: see [`Self::embed_document_with_model_outcome`] — the same
423    /// incomparability applies to batch-indexed vectors when switching models.
424    ///
425    /// Returns `UnknownModel` if `model_name` is not registered.
426    /// Returns an error when any input is bounded; use the outcomes method to
427    /// retain the bounded vectors and per-document byte counts.
428    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            // Atomic preparation records issued work even if its task is cancelled.
488            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    /// Embed a batch of documents for indexing using the configured default model.
529    ///
530    /// Convenience delegate to [`Self::embed_document_batch_with_model`]. Use for
531    /// bulk knowledge-atom and section indexing paths.
532    ///
533    /// Returns `Unconfigured("embedding_model")` if no model is configured.
534    /// Returns an error when any input is bounded; use
535    /// [`Self::embed_document_batch_outcomes`] to retain the bounded vectors.
536    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    /// Embed documents with the default model and retain actual input-bounding metadata.
546    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    /// Embed a batch of queries for retrieval using the named model.
559    ///
560    /// Applies the same pre-0.6 query-role compatibility behavior as
561    /// [`Self::embed_query_with_model`].
562    ///
563    /// Returns `UnknownModel` if `model_name` is not registered.
564    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    /// The note-search vector leg uses the pack-owned graph when that model
588    /// has an installed, consumer-protected bridge. Only a missing graph for
589    /// this consumer takes the existing exact sqlite-vec route.
590    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    /// Hybrid search: text (FTS5) + vector retrieval fused via Reciprocal Rank Fusion.
623    ///
624    /// - Always performs text search over `query_text`.
625    /// - If `query_vector` is `Some`, also performs vector search and fuses both lists.
626    /// - If `None`, returns text-only results — no vector store needed.
627    /// - If `entity_kind` is `Some`, the alive-set query filters to that kind.
628    ///   The text/vector candidate pools are unfiltered up front; the kind
629    ///   filter applies at the alive-check stage where we already fetch each
630    ///   candidate to confirm it isn't soft-deleted.
631    /// - `tags_any`: when non-empty, only entities that have at least one of these
632    ///   tags (case-insensitive) survive the alive-set intersection. Applied BEFORE
633    ///   truncation so matches ranked beyond `limit` in the raw fusion are not lost.
634    /// - `properties_filter`: when `Some`, only entities whose properties are a
635    ///   superset of the given JSON object survive. Applied BEFORE truncation.
636    ///
637    /// `limit` caps the final returned list; internally pulls `limit * 4` candidates per path.
638    ///
639    /// # Cross-namespace visibility (entity search — text leg: visible set; vector leg: primary)
640    ///
641    /// The two legs of entity search scope namespaces differently.
642    ///
643    /// **Text leg.** Entity full-text search is one shared table (`fts_entities`)
644    /// with a `namespace` column. The token's whole visible set (`visible_ns`) is
645    /// forwarded in `TextFilter.namespaces`, which the store applies as a
646    /// `namespace IN (...)` predicate, so one query returns text hits from every
647    /// visible namespace.
648    ///
649    /// **Vector leg.** The vector leg runs only when a query vector is supplied or
650    /// an embedding model is configured. It is an exact sqlite-vec search
651    /// (brute-force cosine, as in `knn`); no ANN index is involved. It issues a
652    /// single request scoped to the primary namespace, because a
653    /// `VectorSearchRequest` carries one namespace. Searching the visible set
654    /// would take one request per namespace and a merge of the per-namespace
655    /// lists; that fanout is not implemented for entities, so entity vector hits
656    /// come from the primary namespace only.
657    ///
658    /// The fused candidates are then checked against the entity store with the
659    /// visible set, so an entity from an extra visible namespace is returned when
660    /// the text leg matched it, and is not returned on a vector-only match.
661    ///
662    /// This differs from `memory.recall`, whose vector leg already searches every
663    /// visible namespace.
664    ///
665    /// Callers can also read any visible entity directly via `get_entity` /
666    /// `resolve`.
667    #[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    /// Hybrid search with an explicit lexical mode for the text arm.
694    #[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    /// Hybrid search over several entity kinds that issues the vector query once.
727    ///
728    /// Returns one list per entry of `entity_kinds`, in that order. Each list is what
729    /// [`Self::hybrid_search`] returns for that kind with the same `limit`: the text stage
730    /// is filtered to the kind and keeps its own `limit * 4` budget, and fusion, the kind
731    /// filter and the cut to `limit` run per kind. Only the vector stage is shared, because
732    /// it takes no entity kind: the single query is the one every per-kind search would have
733    /// issued. When `query_vector` is `None` and an embedding model is configured, the query
734    /// text is embedded once instead of once per kind. No entity-type, tag or property
735    /// filter is set. An empty `entity_kinds` returns no lists and runs no query.
736    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    /// `vector_similarity_floor` is a cosine-similarity value in `[-1.0,
775    /// 1.0]`, matching both the canonical vector-store score contract and the
776    /// scale documented on the `resolve` verb and on
777    /// `SEARCH_VECTOR_SIMILARITY_FLOOR`.
778    #[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    /// Coordinator fan-out variant of [`Self::hybrid_search`]: the text arm
811    /// still fails loud (propagated through `hybrid_search_inner`'s
812    /// `tolerate_vector_error=false` semantics for that leg), but an
813    /// ordinary vector-arm failure after a successful text leg is captured
814    /// instead of discarding the text hits. Request cancellation and the outer
815    /// deadline remain fatal. Reserved for
816    /// `SubstrateCoordinator::fan_out_search_with_visibility` — every other
817    /// caller keeps the fail-loud [`Self::hybrid_search`] contract, so a
818    /// single backend's vector-store outage does not misreport that
819    /// backend's text arm as failed too.
820    #[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    /// Coordinator variant with an explicit lexical mode for the text arm.
845    #[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    /// `vector_pool`, when `Some`, supplies the vector stage's hits and the stage does not
877    /// run; `None` runs it. The stage takes no entity kind, so one pool serves every kind.
878    #[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            // sanitize_fts5_query strips known-unsafe FTS5 metacharacters up front, but if
903            // the lexical leg still errors at runtime on residual punctuation the sanitizer
904            // doesn't strip, this fails loud instead of degrading to vector-only fusion.
905            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                            // Push the entity-kind filter into the FTS query. Without it the
917                            // text arm returns the top `candidates` rows across EVERY entity
918                            // kind in the namespace and the kind is applied only afterwards,
919                            // so when one kind dominates the lexical ranking a search for a
920                            // rarer kind gets back fewer rows than exist, or none. The
921                            // `EntityFilter.kinds` check below stays as the backstop.
922                            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            // The stages read nothing from each other, so they run together; a text error wins.
934            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            // FtsPasses is counted inside the store's `search()` (khive-db
955            // stores/text.rs), only once a real FTS5 statement is prepared —
956            // an empty/fully-sanitized query short-circuits there before any
957            // statement exists and must not count.
958            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            // Each arm fetched `candidates` independently, so their union can contain
966            // twice that many distinct IDs. Keep the complete fetched pool through
967            // ranking and the alive/kind/tag/property filters below; reusing one arm's
968            // cap here lets stale or filtered hits hide live candidates from the other.
969            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            // tags_any has a SQL column and is pushed into query_entities; properties
978            // filtering has no SQL column and is applied in Rust below on the fetched records.
979            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                    // Drop non-matching candidates here, before the alive set is built,
1009                    // so they're excluded ahead of truncation.
1010                    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                // Enrich vector-only hits (title/snippet == None) from entity record.
1022                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    /// Backfill vector and FTS index entries for entities and notes that are missing them.
1082    ///
1083    /// Intended to run once at startup as a background task (warm-up sequence steps 2–4).
1084    /// Queries the SQL substrate for entity bodies and note contents that have no
1085    /// corresponding entry in an eligible embedding model's vector store, then
1086    /// embeds and inserts them. FTS entries missing for notes are also repopulated.
1087    ///
1088    /// The operation is best-effort: individual embed/insert failures are logged and
1089    /// skipped rather than aborting the whole backfill. If no embedding models are
1090    /// registered, returns immediately with 0.
1091    ///
1092    /// Returns the total number of records backfilled across all models.
1093    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            // Must match vec_model_key's naming logic.
1118            let vec_table = format!("vec_{}", sanitize_key(model_name));
1119
1120            // --- Entities: embed the canonical body where no vector exists ---
1121            // Keyset pagination advances past failed or ineligible records too;
1122            // they must not keep the first page full forever.
1123            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            // --- Notes: embed content where no vector entry exists ---
1230            let text_store = self.text_for_notes(token).ok();
1231            let note_store = self.notes(token).ok();
1232            let mut note_total = 0usize;
1233            // Excluded kinds remain absent from this model's vector table.
1234            // Advancing by id prevents a full page of them from repeating.
1235            let mut note_cursor = String::new();
1236            loop {
1237                // Only the id is selected here; the full Note is fetched below so
1238                // note_fts_document gets all fields and stays parity-correct.
1239                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 &note_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 &note_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                    // Repopulate FTS entry using the shared constructor (first model only
1326                    // to avoid N identical overwrites per note).
1327                    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(&note)).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(&note.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, &note, 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    /// Sweep orphaned vector entries for all registered embedding models.
1390    ///
1391    /// A vector entry is orphaned when its `subject_id` no longer exists as a
1392    /// live row in the entity, note, or knowledge-atom tables (i.e. either the
1393    /// row is absent or has `deleted_at IS NOT NULL`). Orphaned entries
1394    /// accumulate after hard-deletes because the vector store and SQL
1395    /// substrate are decoupled.
1396    ///
1397    /// Iterates over every registered embedding model and calls
1398    /// [`khive_storage::VectorStore::orphan_sweep`] for the token's namespace. Models whose
1399    /// backend returns [`khive_storage::StorageError::Unsupported`] are skipped without error —
1400    /// this preserves forward-compat when a newly registered model does not yet
1401    /// implement sweep.
1402    ///
1403    /// Returns the total number of vector rows deleted across all models.
1404    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
1492/// Returns `true` when `properties` is a superset of all key-value pairs in `filter`.
1493///
1494/// Mirrors the semantics of `khive_pack_kg::handlers::common::props_match` so that the
1495/// storage-leg predicate is identical to the handler-side post-filter.
1496pub(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
1513/// Score bonus applied when an entity's title is an exact case-insensitive match for
1514/// the query. Dominates RRF scores (~0.09–0.18 range with k=10) so that an exact
1515/// name match always ranks above any partial or semantic match.
1516const EXACT_MATCH_BOOST: f64 = 0.5;
1517
1518/// Fuse text + vector hits with Reciprocal Rank Fusion (k=10).
1519///
1520/// Scoring and the merge of each id's labels come from the shared labelled fusion in
1521/// `khive-retrieval`. Entity search keeps its own k=10 and exact-match boosting here.
1522/// Hits in both lists get RRF scores summed. If `query_text` exactly matches
1523/// (case-insensitive) the title a fused hit carries from the text hits, a bonus of
1524/// `EXACT_MATCH_BOOST` is added to ensure exact-name matches dominate.
1525/// Sort by fused score, take top-`limit`.
1526fn 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    // The bonus joins the full fused order, so it decides the sort and the cut below.
1568    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;