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 uuid::Uuid;
6
7use crate::config::{parse_embedding_model_alias, sanitize_key};
8use crate::curation::note_fts_document;
9use crate::error::{RuntimeError, RuntimeResult};
10use crate::runtime::{KhiveRuntime, NamespaceToken};
11use khive_score::{rrf_score, DeterministicScore};
12use khive_storage::types::{
13    PageRequest, TextFilter, TextQueryMode, TextSearchHit, TextSearchRequest, VectorSearchHit,
14    VectorSearchRequest,
15};
16use khive_storage::EntityFilter;
17use khive_types::SubstrateKind;
18
19// Fault-injection point for backfill reader errors. When set, the next
20// `backfill_missing_embeddings` call substitutes a `StorageError::Pool` for the
21// result of `sql.reader().await`, then resets to false. Available in test builds
22// and the `fault-injection` feature for integration testing.
23#[cfg(any(test, feature = "fault-injection"))]
24std::thread_local! {
25    static BACKFILL_READER_FAIL: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
26}
27
28/// Arm the backfill reader fault injection. When set, the next call to
29/// `backfill_missing_embeddings` will substitute a `StorageError::Pool` for the
30/// result of `sql.reader().await`, then reset the flag. The injected error
31/// passes through `map_err(RuntimeError::Storage)?` — the same path as a real
32/// reader failure — so it exercises the fail-closed guard rather than bypassing it.
33/// Available when compiled with `cfg(test)` or `feature = "fault-injection"`.
34#[cfg(any(test, feature = "fault-injection"))]
35pub fn arm_backfill_reader_fail() {
36    BACKFILL_READER_FAIL.with(|c| c.set(true));
37}
38
39/// A unified search result combining vector and text signals.
40#[derive(Clone, Debug)]
41pub struct SearchHit {
42    pub entity_id: Uuid,
43    pub score: DeterministicScore,
44    pub source: SearchSource,
45    pub title: Option<String>,
46    pub snippet: Option<String>,
47}
48
49/// Which retrieval path(s) contributed to a hit.
50#[derive(Clone, Copy, Debug, PartialEq, Eq)]
51pub enum SearchSource {
52    Vector,
53    Text,
54    Both,
55}
56
57/// RRF constant. Controls how strongly top ranks dominate.
58///
59/// The original paper uses k=60 for large-scale document retrieval. For a knowledge
60/// graph with tens to thousands of entities, k=60 over-compresses scores into a
61/// narrow band (rank 1 ≈ 0.016, rank 10 ≈ 0.014, spread ≈ 0.002). k=10 produces
62/// rank 1 ≈ 0.091, rank 10 ≈ 0.050, spread ≈ 0.041 — 20× better discrimination,
63/// making dedup-before-create reliable 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
69impl KhiveRuntime {
70    /// Generate an embedding vector for `text` using the configured default model.
71    ///
72    /// First call lazily loads model weights (cold start cost). Subsequent calls reuse them.
73    /// Returns `Unconfigured("embedding_model")` if no model is configured.
74    pub async fn embed(&self, text: &str) -> RuntimeResult<Vec<f32>> {
75        let model_name = self.default_embedder_name();
76        if model_name.is_empty() {
77            return Err(RuntimeError::Unconfigured("embedding_model".into()));
78        }
79        self.embed_with_model(model_name, text).await
80    }
81
82    /// Generate an embedding vector for `text` using the named model.
83    ///
84    /// Accepts both built-in lattice model names/aliases and custom provider
85    /// names registered via [`KhiveRuntime::register_embedder`]. For lattice
86    /// models the resolved `EmbeddingModel` enum is forwarded to `embed_one`
87    /// so the service can select the correct model variant. For custom
88    /// providers, `embed_one` is called with `EmbeddingModel::default()`
89    /// because custom services are expected to ignore the enum argument (they
90    /// own a single model implicitly).
91    ///
92    /// Applies no instruction prefix (generic role). Use
93    /// [`Self::embed_document_with_model`] / [`Self::embed_query_with_model`] for
94    /// instruction-tuned models where the asymmetric prefix matters.
95    ///
96    /// Returns `UnknownModel` if `model_name` is not in the embedder registry.
97    pub async fn embed_with_model(&self, model_name: &str, text: &str) -> RuntimeResult<Vec<f32>> {
98        // Try to resolve as a lattice alias. If that succeeds, use the enum to
99        // inform the service which model to run. If not, fall through to the
100        // custom-provider path — custom services ignore the EmbeddingModel arg.
101        let model = parse_embedding_model_alias(model_name);
102        let service = self.embedder(model_name).await?;
103        let emb_model = model.unwrap_or_default();
104        Ok(service.embed_one(text, emb_model).await?)
105    }
106
107    /// Embed a document/passage for indexing using the named model.
108    ///
109    /// Applies `EmbeddingService::embed_passage`, which prepends the model's
110    /// `document_instruction()` prefix when defined (e.g. `"passage: "` for
111    /// multilingual-e5). For models with no document prefix (MiniLM, BGE) this
112    /// is identical to [`Self::embed_with_model`].
113    ///
114    /// Use this for all index/store/backfill paths so that instruction-tuned
115    /// models produce passage-side vectors.
116    ///
117    /// **Reindex caveat**: switching from an unprefixed model (or a model with no
118    /// `document_instruction`) to an instruction-tuned model changes the vector
119    /// representation. Vectors stored under the old scheme are not comparable to
120    /// newly prefixed vectors. Operators must trigger a full reindex
121    /// (`knowledge.index(rebuild_ann=true)` / `kkernel reindex`) after changing
122    /// the embedding model config.
123    ///
124    /// Returns `UnknownModel` if `model_name` is not registered.
125    pub async fn embed_document_with_model(
126        &self,
127        model_name: &str,
128        text: &str,
129    ) -> RuntimeResult<Vec<f32>> {
130        let model = parse_embedding_model_alias(model_name);
131        let service = self.embedder(model_name).await?;
132        let emb_model = model.unwrap_or_default();
133        service
134            .embed_passage(&[text.to_string()], emb_model)
135            .await?
136            .into_iter()
137            .next()
138            .ok_or_else(|| RuntimeError::Internal("embed_passage returned empty vec".into()))
139    }
140
141    /// Embed a query string for retrieval using the named model.
142    ///
143    /// Applies `EmbeddingService::embed_query`, which prepends the model's
144    /// `query_instruction()` prefix when defined (e.g. `"query: "` for
145    /// multilingual-e5). For models with no query prefix (MiniLM, BGE) this
146    /// is identical to [`Self::embed_with_model`].
147    ///
148    /// Use this for all search/recall/suggest query embedding paths so that
149    /// instruction-tuned models land in the correct side of their retrieval
150    /// space.
151    ///
152    /// Returns `UnknownModel` if `model_name` is not registered.
153    pub async fn embed_query_with_model(
154        &self,
155        model_name: &str,
156        text: &str,
157    ) -> RuntimeResult<Vec<f32>> {
158        let model = parse_embedding_model_alias(model_name);
159        let service = self.embedder(model_name).await?;
160        let emb_model = model.unwrap_or_default();
161        service
162            .embed_query(&[text.to_string()], emb_model)
163            .await?
164            .into_iter()
165            .next()
166            .ok_or_else(|| RuntimeError::Internal("embed_query returned empty vec".into()))
167    }
168
169    /// Embed a document for indexing using the configured default model.
170    ///
171    /// Delegates to [`Self::embed_document_with_model`]. Use for entity/note
172    /// create and reindex paths.
173    ///
174    /// Returns `Unconfigured("embedding_model")` if no model is configured.
175    pub async fn embed_document(&self, text: &str) -> RuntimeResult<Vec<f32>> {
176        let model_name = self.default_embedder_name();
177        if model_name.is_empty() {
178            return Err(RuntimeError::Unconfigured("embedding_model".into()));
179        }
180        self.embed_document_with_model(model_name, text).await
181    }
182
183    /// Embed a query for retrieval using the configured default model.
184    ///
185    /// Delegates to [`Self::embed_query_with_model`]. Use for vector search and
186    /// hybrid search query paths.
187    ///
188    /// Returns `Unconfigured("embedding_model")` if no model is configured.
189    pub async fn embed_query(&self, text: &str) -> RuntimeResult<Vec<f32>> {
190        let model_name = self.default_embedder_name();
191        if model_name.is_empty() {
192            return Err(RuntimeError::Unconfigured("embedding_model".into()));
193        }
194        self.embed_query_with_model(model_name, text).await
195    }
196
197    /// Generate embeddings for multiple texts in one call using the configured default model.
198    ///
199    /// Delegates to the cached `EmbeddingService::embed`, so repeated texts within
200    /// and across calls benefit from the runtime-level LRU cache.
201    ///
202    /// Returns an empty vec for empty input without hitting the embedding service.
203    /// Returns `Unconfigured("embedding_model")` if no model is configured.
204    pub async fn embed_batch(&self, texts: &[String]) -> RuntimeResult<Vec<Vec<f32>>> {
205        if texts.is_empty() {
206            return Ok(vec![]);
207        }
208        let model_name = self.default_embedder_name();
209        if model_name.is_empty() {
210            return Err(RuntimeError::Unconfigured("embedding_model".into()));
211        }
212        self.embed_batch_with_model(model_name, texts).await
213    }
214
215    /// Generate embeddings for multiple texts using the named model.
216    ///
217    /// Accepts lattice model names/aliases and custom provider names.
218    /// Returns `UnknownModel` if `model_name` is not in the embedder registry.
219    pub async fn embed_batch_with_model(
220        &self,
221        model_name: &str,
222        texts: &[String],
223    ) -> RuntimeResult<Vec<Vec<f32>>> {
224        if texts.is_empty() {
225            return Ok(vec![]);
226        }
227        let model = parse_embedding_model_alias(model_name);
228        let service = self.embedder(model_name).await?;
229        let emb_model = model.unwrap_or_default();
230        Ok(service.embed(texts, emb_model).await?)
231    }
232
233    /// Embed a batch of documents for indexing using the named model.
234    ///
235    /// Applies `EmbeddingService::embed_passage`. Use for all bulk
236    /// index/backfill/reindex operations to apply the passage-side prefix.
237    ///
238    /// **Reindex caveat**: see [`Self::embed_document_with_model`] — the same
239    /// incomparability applies to batch-indexed vectors when switching models.
240    ///
241    /// Returns `UnknownModel` if `model_name` is not registered.
242    pub async fn embed_document_batch_with_model(
243        &self,
244        model_name: &str,
245        texts: &[String],
246    ) -> RuntimeResult<Vec<Vec<f32>>> {
247        if texts.is_empty() {
248            return Ok(vec![]);
249        }
250        let model = parse_embedding_model_alias(model_name);
251        let service = self.embedder(model_name).await?;
252        let emb_model = model.unwrap_or_default();
253        Ok(service.embed_passage(texts, emb_model).await?)
254    }
255
256    /// Embed a batch of documents for indexing using the configured default model.
257    ///
258    /// Convenience delegate to [`Self::embed_document_batch_with_model`]. Use for
259    /// bulk knowledge-atom and section indexing paths.
260    ///
261    /// Returns `Unconfigured("embedding_model")` if no model is configured.
262    pub async fn embed_document_batch(&self, texts: &[String]) -> RuntimeResult<Vec<Vec<f32>>> {
263        if texts.is_empty() {
264            return Ok(vec![]);
265        }
266        let model_name = self.default_embedder_name();
267        if model_name.is_empty() {
268            return Err(RuntimeError::Unconfigured("embedding_model".into()));
269        }
270        self.embed_document_batch_with_model(model_name, texts)
271            .await
272    }
273
274    /// Embed a batch of queries for retrieval using the named model.
275    ///
276    /// Applies `EmbeddingService::embed_query`. Use for bulk query-side
277    /// operations where multiple queries need instruction-tuned prefixing.
278    ///
279    /// Returns `UnknownModel` if `model_name` is not registered.
280    pub async fn embed_query_batch_with_model(
281        &self,
282        model_name: &str,
283        texts: &[String],
284    ) -> RuntimeResult<Vec<Vec<f32>>> {
285        if texts.is_empty() {
286            return Ok(vec![]);
287        }
288        let model = parse_embedding_model_alias(model_name);
289        let service = self.embedder(model_name).await?;
290        let emb_model = model.unwrap_or_default();
291        Ok(service.embed_query(texts, emb_model).await?)
292    }
293
294    /// Search vectors using either a caller-provided embedding or query text.
295    ///
296    /// Existing callers pass `query_embedding: Some(vec)` to avoid re-embedding.
297    /// Text callers pass `query_embedding: None, query_text: Some(...)` and the
298    /// runtime embeds internally.
299    pub async fn vector_search(
300        &self,
301        token: &NamespaceToken,
302        query_embedding: Option<Vec<f32>>,
303        query_text: Option<&str>,
304        top_k: u32,
305        kind: Option<SubstrateKind>,
306    ) -> RuntimeResult<Vec<VectorSearchHit>> {
307        let embedding = match query_embedding {
308            Some(vec) => vec,
309            None => {
310                let text = query_text.ok_or_else(|| {
311                    RuntimeError::InvalidInput(
312                        "vector search requires query_embedding or query_text".into(),
313                    )
314                })?;
315                if text.trim().is_empty() {
316                    return Err(RuntimeError::InvalidInput(
317                        "query_text must not be empty".into(),
318                    ));
319                }
320                self.embed_query(text).await?
321            }
322        };
323
324        let ns = token.namespace().as_str().to_owned();
325        Ok(self
326            .vectors(token)?
327            .search(VectorSearchRequest {
328                query_vectors: vec![embedding],
329                top_k,
330                namespace: Some(ns),
331                kind,
332                embedding_model: None,
333                filter: None,
334                backend_hints: None,
335            })
336            .await?)
337    }
338
339    /// Hybrid search: text (FTS5) + vector retrieval fused via Reciprocal Rank Fusion.
340    ///
341    /// - Always performs text search over `query_text`.
342    /// - If `query_vector` is `Some`, also performs vector search and fuses both lists.
343    /// - If `None`, returns text-only results — no vector store needed.
344    /// - If `entity_kind` is `Some`, the alive-set query filters to that kind.
345    ///   The text/vector candidate pools are unfiltered up front; the kind
346    ///   filter applies at the alive-check stage where we already fetch each
347    ///   candidate to confirm it isn't soft-deleted.
348    /// - `tags_any`: when non-empty, only entities that have at least one of these
349    ///   tags (case-insensitive) survive the alive-set intersection. Applied BEFORE
350    ///   truncation so matches ranked beyond `limit` in the raw fusion are not lost.
351    /// - `properties_filter`: when `Some`, only entities whose properties are a
352    ///   superset of the given JSON object survive. Applied BEFORE truncation.
353    ///
354    /// `limit` caps the final returned list; internally pulls `limit * 4` candidates per path.
355    /// The fused candidate set is kept untruncated until after the alive + kind + tag +
356    /// properties filter so that matching hits ranked below `limit` in the raw fusion
357    /// still surface when higher-ranked candidates are excluded by any filter.
358    ///
359    /// # Cross-namespace visibility (entity search — primary namespace only; deferred)
360    ///
361    /// Both the **FTS leg** and the **vector/ANN leg** of entity search (`hybrid_search`)
362    /// are restricted to the **primary namespace only**.
363    ///
364    /// Rationale: each namespace owns a separate FTS table (`fts_entities_{ns}`)
365    /// and a separate ANN index instance. Cross-namespace entity-search fanout
366    /// requires iterating over every visible namespace's store, issuing parallel
367    /// search requests, and fusing the results — this is deferred (entity-search
368    /// cross-namespace fanout is the outstanding follow-up).
369    ///
370    /// Note: this is distinct from memory recall's cross-namespace fanout, which
371    /// ships in ADR-007 Rev 4 (`memory.recall` iterates `visible_namespaces` across
372    /// both the FTS and vector legs). Entity search fanout is the remaining deferred
373    /// piece; memory recall fanout is not deferred.
374    ///
375    /// The `visible_ns` list is forwarded in the `TextFilter.namespaces` field,
376    /// which limits results to those namespaces within the primary store. Because
377    /// entities from extra namespaces live in their own FTS tables, this filter has
378    /// no cross-namespace effect today.
379    ///
380    /// Callers with a multi-namespace visible set can READ cross-namespace entities
381    /// directly via `get_entity` / `resolve`, but `hybrid_search` returns only
382    /// primary-namespace hits until entity-search cross-namespace fanout ships.
383    #[allow(clippy::too_many_arguments)]
384    pub async fn hybrid_search(
385        &self,
386        token: &NamespaceToken,
387        query_text: &str,
388        query_vector: Option<Vec<f32>>,
389        limit: u32,
390        entity_kind: Option<&str>,
391        entity_type: Option<&str>,
392        tags_any: &[String],
393        properties_filter: Option<&serde_json::Value>,
394    ) -> RuntimeResult<Vec<SearchHit>> {
395        let candidates = limit.saturating_mul(CANDIDATE_MULTIPLIER).max(limit);
396
397        let visible_ns: Vec<String> = token
398            .visible_namespaces()
399            .iter()
400            .map(|ns| ns.as_str().to_owned())
401            .collect();
402        let text_hits = self
403            .text(token)?
404            .search(TextSearchRequest {
405                query: query_text.to_string(),
406                mode: TextQueryMode::Plain,
407                filter: Some(TextFilter {
408                    namespaces: visible_ns.clone(),
409                    ..TextFilter::default()
410                }),
411                top_k: candidates,
412                snippet_chars: 200,
413            })
414            .await?;
415
416        let vector_hits = if query_vector.is_some() || self.config().embedding_model.is_some() {
417            self.vector_search(
418                token,
419                query_vector,
420                Some(query_text),
421                candidates,
422                Some(SubstrateKind::Entity),
423            )
424            .await?
425        } else {
426            Vec::new()
427        };
428
429        // Fuse without truncating: keep the full candidate pool through the
430        // alive/kind/tag/property filter so matching hits below rank `limit`
431        // aren't lost when higher-ranked candidates are excluded.
432        let mut fused = rrf_fuse(text_hits, vector_hits, candidates as usize, query_text);
433
434        // Filter to alive entities (and optionally to a specific kind, tags, or
435        // properties). A single query fetches all alive IDs that match the kind
436        // and tag constraints from the fused set; any ID absent has been
437        // soft-deleted or doesn't match. The SQL-level `tags_any` filter is
438        // pushed into `query_entities`; properties filtering (no SQL column)
439        // is applied at the Rust level using the entity records already fetched.
440        if !fused.is_empty() {
441            let candidate_ids: Vec<Uuid> = fused.iter().map(|h| h.entity_id).collect();
442            let alive_page = self
443                .entities(token)?
444                .query_entities(
445                    token.namespace().as_str(),
446                    EntityFilter {
447                        ids: candidate_ids,
448                        kinds: entity_kind.map(|k| vec![k.to_string()]).unwrap_or_default(),
449                        entity_types: entity_type.map(|t| vec![t.to_string()]).unwrap_or_default(),
450                        namespaces: visible_ns,
451                        tags_any: tags_any.to_vec(),
452                        ..EntityFilter::default()
453                    },
454                    PageRequest {
455                        offset: 0,
456                        limit: fused.len() as u32,
457                    },
458                )
459                .await?;
460            // Keep entity metadata to enrich hits that had no FTS5 title/snippet,
461            // and to apply the properties filter before truncation.
462            let mut entity_meta: HashMap<Uuid, (String, Option<String>)> = HashMap::new();
463            let mut alive: HashSet<Uuid> = HashSet::new();
464            for e in alive_page.items {
465                // Apply properties predicate here — before adding to the alive set —
466                // so that non-matching candidates are dropped before truncation.
467                if let Some(pf) = properties_filter {
468                    if !entity_props_match(e.properties.as_ref(), pf) {
469                        continue;
470                    }
471                }
472                alive.insert(e.id);
473                entity_meta.insert(e.id, (e.name, e.description));
474            }
475
476            fused.retain(|h| alive.contains(&h.entity_id));
477
478            // Enrich vector-only hits (title/snippet == None) from entity record.
479            for hit in &mut fused {
480                if let Some((name, description)) = entity_meta.get(&hit.entity_id) {
481                    if hit.title.is_none() {
482                        hit.title = Some(name.clone());
483                    }
484                    if hit.snippet.is_none() {
485                        hit.snippet = description.clone();
486                    }
487                }
488            }
489        }
490
491        fused.truncate(limit as usize);
492        Ok(fused)
493    }
494
495    /// Exact KNN over the full namespace's vector store.
496    ///
497    /// sqlite-vec uses brute-force cosine — results are exact, not approximate.
498    /// Cost is O(N · D) per query. For small-to-medium namespaces (~hundreds of
499    /// thousands of vectors) this is well within latency budgets.
500    pub async fn knn(
501        &self,
502        token: &NamespaceToken,
503        query_vector: Vec<f32>,
504        top_k: u32,
505    ) -> RuntimeResult<Vec<VectorSearchHit>> {
506        let ns = token.namespace().as_str().to_owned();
507        Ok(self
508            .vectors(token)?
509            .search(VectorSearchRequest {
510                query_vectors: vec![query_vector],
511                top_k,
512                namespace: Some(ns),
513                kind: Some(SubstrateKind::Entity),
514                embedding_model: None,
515                filter: None,
516                backend_hints: None,
517            })
518            .await?)
519    }
520
521    /// Exact KNN restricted to a candidate set.
522    ///
523    /// Useful for reranking the top-N results from `hybrid_search` (or any other
524    /// retrieval path) with exact cosine similarity against a query vector.
525    /// Returns hits sorted by similarity (highest first), truncated to `top_k`.
526    pub async fn rerank(
527        &self,
528        token: &NamespaceToken,
529        query_vector: &[f32],
530        candidate_ids: &[Uuid],
531        top_k: u32,
532    ) -> RuntimeResult<Vec<VectorSearchHit>> {
533        let candidate_set: HashSet<Uuid> = candidate_ids.iter().copied().collect();
534        let ns = token.namespace().as_str().to_owned();
535        let all_hits = self
536            .vectors(token)?
537            .search(VectorSearchRequest {
538                query_vectors: vec![query_vector.to_vec()],
539                top_k: candidate_ids.len() as u32,
540                namespace: Some(ns),
541                kind: Some(SubstrateKind::Entity),
542                embedding_model: None,
543                filter: None,
544                backend_hints: None,
545            })
546            .await?;
547        let mut hits: Vec<VectorSearchHit> = all_hits
548            .into_iter()
549            .filter(|h| candidate_set.contains(&h.subject_id))
550            .collect();
551        hits.sort_by(|a, b| b.score.cmp(&a.score));
552        hits.truncate(top_k as usize);
553        Ok(hits)
554    }
555
556    /// Backfill vector and FTS index entries for entities and notes that are missing them.
557    ///
558    /// Intended to run once at startup as a background task (warm-up sequence steps 2–4).
559    /// Queries the SQL substrate for entity descriptions and note contents that have no
560    /// corresponding entry in the vector store for any registered embedding model, then
561    /// embeds and inserts them. FTS entries missing for notes are also repopulated.
562    ///
563    /// The operation is best-effort: individual embed/insert failures are logged and
564    /// skipped rather than aborting the whole backfill. If no embedding models are
565    /// registered, returns immediately with 0.
566    ///
567    /// Returns the total number of records backfilled across all models.
568    pub async fn backfill_missing_embeddings(&self, token: &NamespaceToken) -> RuntimeResult<u64> {
569        use khive_storage::types::{SqlRow, SqlStatement, SqlValue};
570
571        let model_names = self.registered_embedding_model_names();
572        if model_names.is_empty() {
573            tracing::debug!(
574                "backfill_missing_embeddings: no embedding models registered, skipping"
575            );
576            return Ok(0);
577        }
578
579        let ns = token.namespace().as_str().to_string();
580        let mut total_backfilled = 0u64;
581
582        for model_name in &model_names {
583            // Derive the vec table name from the model name (must match vec_model_key logic).
584            let vec_table = format!("vec_{}", sanitize_key(model_name));
585
586            // --- Entities: embed description where no vector entry exists ---
587            // Loop until a batch returns fewer than PAGE_SIZE rows. Because the query uses
588            // NOT IN (SELECT subject_id FROM vec_table ...), each successfully inserted row is
589            // excluded from subsequent pages — no OFFSET needed.
590            const PAGE_SIZE: usize = 500;
591            let mut entity_total = 0usize;
592            loop {
593                let entity_sql = SqlStatement {
594                    sql: format!(
595                        "SELECT id, name, description FROM entities \
596                         WHERE namespace = ?1 AND deleted_at IS NULL \
597                         AND id NOT IN (\
598                             SELECT subject_id FROM {vec_table} \
599                             WHERE namespace = ?1 AND embedding_model = ?2 \
600                         ) LIMIT {PAGE_SIZE}"
601                    ),
602                    params: vec![
603                        SqlValue::Text(ns.clone()),
604                        SqlValue::Text(model_name.clone()),
605                    ],
606                    label: Some("backfill_entities".into()),
607                };
608
609                let entity_rows: Vec<SqlRow> = {
610                    let sql = self.sql();
611                    let reader_result = sql.reader().await;
612                    #[cfg(any(test, feature = "fault-injection"))]
613                    let reader_result = if BACKFILL_READER_FAIL.with(|c| c.get()) {
614                        BACKFILL_READER_FAIL.with(|c| c.set(false));
615                        Err(khive_storage::StorageError::Pool {
616                            operation: "reader".into(),
617                            message: "injected failure".into(),
618                        })
619                    } else {
620                        reader_result
621                    };
622                    let mut reader = reader_result.map_err(RuntimeError::Storage)?;
623                    reader
624                        .query_all(entity_sql)
625                        .await
626                        .map_err(RuntimeError::Storage)?
627                };
628
629                let batch_len = entity_rows.len();
630                entity_total += batch_len;
631
632                for row in &entity_rows {
633                    let id_str = row.columns.first().and_then(|c| {
634                        if let SqlValue::Text(s) = &c.value {
635                            Some(s.clone())
636                        } else {
637                            None
638                        }
639                    });
640                    let description = row.columns.get(2).and_then(|c| {
641                        if let SqlValue::Text(s) = &c.value {
642                            Some(s.clone())
643                        } else if let SqlValue::Null = &c.value {
644                            None
645                        } else {
646                            None
647                        }
648                    });
649
650                    let (Some(id_str), Some(desc)) = (id_str, description) else {
651                        continue;
652                    };
653                    let Ok(id) = id_str.parse::<Uuid>() else {
654                        continue;
655                    };
656                    if desc.trim().is_empty() {
657                        continue;
658                    }
659
660                    match self.embed_document_with_model(model_name, &desc).await {
661                        Ok(vector) => {
662                            if let Ok(vs) = self.vectors_for_model(token, model_name) {
663                                match vs
664                                    .insert(
665                                        id,
666                                        SubstrateKind::Entity,
667                                        &ns,
668                                        "entity.description",
669                                        vec![vector],
670                                    )
671                                    .await
672                                {
673                                    Ok(()) => {
674                                        total_backfilled += 1;
675                                    }
676                                    Err(e) => {
677                                        tracing::warn!(
678                                            id = %id, model = %model_name,
679                                            error = %e,
680                                            "backfill_missing_embeddings: entity vector insert failed"
681                                        );
682                                    }
683                                }
684                            }
685                        }
686                        Err(e) => {
687                            tracing::warn!(
688                                id = %id, model = %model_name,
689                                error = %e,
690                                "backfill_missing_embeddings: entity embed failed"
691                            );
692                        }
693                    }
694                }
695
696                if batch_len < PAGE_SIZE {
697                    break;
698                }
699            }
700
701            // --- Notes: embed content where no vector entry exists ---
702            let text_store = self.text_for_notes(token).ok();
703            let note_store = self.notes(token).ok();
704            let mut note_total = 0usize;
705            loop {
706                // Select only the id here; the full Note is fetched below so that
707                // note_fts_document receives all fields (name, properties, updated_at)
708                // and produces a parity-correct document rather than a stripped one.
709                let note_sql = SqlStatement {
710                    sql: format!(
711                        "SELECT id FROM notes \
712                         WHERE namespace = ?1 AND deleted_at IS NULL \
713                         AND id NOT IN (\
714                             SELECT subject_id FROM {vec_table} \
715                             WHERE namespace = ?1 AND embedding_model = ?2 \
716                         ) LIMIT {PAGE_SIZE}"
717                    ),
718                    params: vec![
719                        SqlValue::Text(ns.clone()),
720                        SqlValue::Text(model_name.clone()),
721                    ],
722                    label: Some("backfill_notes".into()),
723                };
724
725                let note_rows: Vec<SqlRow> = {
726                    let sql = self.sql();
727                    let reader_result = sql.reader().await;
728                    #[cfg(any(test, feature = "fault-injection"))]
729                    let reader_result = if BACKFILL_READER_FAIL.with(|c| c.get()) {
730                        BACKFILL_READER_FAIL.with(|c| c.set(false));
731                        Err(khive_storage::StorageError::Pool {
732                            operation: "reader".into(),
733                            message: "injected failure".into(),
734                        })
735                    } else {
736                        reader_result
737                    };
738                    let mut reader = reader_result.map_err(RuntimeError::Storage)?;
739                    reader
740                        .query_all(note_sql)
741                        .await
742                        .map_err(RuntimeError::Storage)?
743                };
744
745                let batch_len = note_rows.len();
746                note_total += batch_len;
747
748                for row in &note_rows {
749                    let id_str = row.columns.first().and_then(|c| {
750                        if let SqlValue::Text(s) = &c.value {
751                            Some(s.clone())
752                        } else {
753                            None
754                        }
755                    });
756
757                    let Some(id_str) = id_str else {
758                        continue;
759                    };
760                    let Ok(id) = id_str.parse::<Uuid>() else {
761                        continue;
762                    };
763
764                    // Fetch the full Note so that note_fts_document has all fields
765                    // (name, properties, updated_at) — prevents overwriting a correct
766                    // FTS row with a stripped content-only document.
767                    let note = match &note_store {
768                        Some(store) => match store.get_note(id).await {
769                            Ok(Some(n)) => n,
770                            _ => continue,
771                        },
772                        None => continue,
773                    };
774
775                    if note.content.trim().is_empty() {
776                        continue;
777                    }
778
779                    // Repopulate FTS entry using the shared constructor (first model only
780                    // to avoid N identical overwrites per note).
781                    if model_names.first().map(|n| n.as_str()) == Some(model_name.as_str()) {
782                        if let Some(ref ts) = text_store {
783                            if let Err(e) = ts.upsert_document(note_fts_document(&note)).await {
784                                tracing::warn!(id = %id, error = %e,
785                                    "backfill_missing_embeddings: note FTS upsert failed");
786                            }
787                        }
788                    }
789
790                    let content = note.content.clone();
791                    match self.embed_document_with_model(model_name, &content).await {
792                        Ok(vector) => {
793                            if let Ok(vs) = self.vectors_for_model(token, model_name) {
794                                match vs
795                                    .insert(
796                                        id,
797                                        SubstrateKind::Note,
798                                        &ns,
799                                        "note.content",
800                                        vec![vector],
801                                    )
802                                    .await
803                                {
804                                    Ok(()) => {
805                                        total_backfilled += 1;
806                                    }
807                                    Err(e) => {
808                                        tracing::warn!(
809                                            id = %id, model = %model_name,
810                                            error = %e,
811                                            "backfill_missing_embeddings: note vector insert failed"
812                                        );
813                                    }
814                                }
815                            }
816                        }
817                        Err(e) => {
818                            tracing::warn!(
819                                id = %id, model = %model_name,
820                                error = %e,
821                                "backfill_missing_embeddings: note embed failed"
822                            );
823                        }
824                    }
825                }
826
827                if batch_len < PAGE_SIZE {
828                    break;
829                }
830            }
831
832            tracing::info!(
833                model = %model_name,
834                namespace = %ns,
835                entities = entity_total,
836                notes = note_total,
837                "backfill_missing_embeddings: model pass complete"
838            );
839        }
840
841        tracing::info!(
842            namespace = %ns,
843            total_backfilled = total_backfilled,
844            "backfill_missing_embeddings: finished"
845        );
846
847        Ok(total_backfilled)
848    }
849
850    /// Sweep orphaned vector entries for all registered embedding models.
851    ///
852    /// A vector entry is orphaned when its `subject_id` no longer exists as a
853    /// live row in the entity or note tables (i.e. either the row is absent or
854    /// has `deleted_at IS NOT NULL`). Orphaned entries accumulate after
855    /// hard-deletes because the vector store and SQL substrate are decoupled.
856    ///
857    /// Iterates over every registered embedding model and calls
858    /// [`khive_storage::VectorStore::orphan_sweep`] for the token's namespace. Models whose
859    /// backend returns [`khive_storage::StorageError::Unsupported`] are skipped without error —
860    /// this preserves forward-compat when a newly registered model does not yet
861    /// implement sweep.
862    ///
863    /// Returns the total number of vector rows deleted across all models.
864    pub async fn sweep_orphan_vectors(
865        &self,
866        token: &NamespaceToken,
867        max_delete_per_model: u32,
868        dry_run: bool,
869    ) -> RuntimeResult<u64> {
870        use khive_storage::types::OrphanSweepConfig;
871        use khive_storage::StorageError;
872
873        let model_names = self.registered_embedding_model_names();
874        if model_names.is_empty() {
875            tracing::debug!("sweep_orphan_vectors: no embedding models registered, skipping");
876            return Ok(0);
877        }
878
879        let ns = token.namespace().as_str().to_string();
880        let mut total_deleted = 0u64;
881
882        for model_name in &model_names {
883            let store = match self.vectors_for_model(token, model_name) {
884                Ok(s) => s,
885                Err(e) => {
886                    tracing::warn!(
887                        model = %model_name,
888                        error = %e,
889                        "sweep_orphan_vectors: failed to get vector store, skipping model"
890                    );
891                    continue;
892                }
893            };
894
895            let caps = store.capabilities();
896            if !caps.supports_orphan_sweep {
897                tracing::debug!(
898                    model = %model_name,
899                    "sweep_orphan_vectors: backend does not support orphan sweep, skipping"
900                );
901                continue;
902            }
903
904            let config = OrphanSweepConfig {
905                subject_id_allowlist: None,
906                namespaces: vec![ns.clone()],
907                substrate_kinds: vec![],
908                max_delete: max_delete_per_model,
909                dry_run,
910            };
911
912            match store.orphan_sweep(&config).await {
913                Ok(result) => {
914                    tracing::info!(
915                        model = %model_name,
916                        namespace = %ns,
917                        scanned = result.scanned,
918                        deleted = result.deleted,
919                        would_delete = result.would_delete,
920                        dry_run = dry_run,
921                        "sweep_orphan_vectors: sweep complete"
922                    );
923                    total_deleted += result.deleted;
924                }
925                Err(StorageError::Unsupported { .. }) => {
926                    tracing::debug!(
927                        model = %model_name,
928                        "sweep_orphan_vectors: backend returned Unsupported, skipping"
929                    );
930                }
931                Err(e) => {
932                    tracing::warn!(
933                        model = %model_name,
934                        error = %e,
935                        "sweep_orphan_vectors: sweep failed, continuing with other models"
936                    );
937                }
938            }
939        }
940
941        tracing::info!(
942            namespace = %ns,
943            total_deleted = total_deleted,
944            dry_run = dry_run,
945            "sweep_orphan_vectors: finished"
946        );
947
948        Ok(total_deleted)
949    }
950}
951
952/// Returns `true` when `entity_props` is a superset of all key-value pairs in `filter`.
953///
954/// Mirrors the semantics of `khive_pack_kg::handlers::common::props_match` so that the
955/// storage-leg predicate is identical to the handler-side post-filter.
956fn entity_props_match(
957    entity_props: Option<&serde_json::Value>,
958    filter: &serde_json::Value,
959) -> bool {
960    let required = match filter.as_object() {
961        Some(obj) if !obj.is_empty() => obj,
962        _ => return true,
963    };
964    let actual = match entity_props.and_then(serde_json::Value::as_object) {
965        Some(obj) => obj,
966        None => return false,
967    };
968    required
969        .iter()
970        .all(|(k, v)| actual.get(k).is_some_and(|av| av == v))
971}
972
973/// Score bonus applied when an entity's title is an exact case-insensitive match for
974/// the query. Dominates RRF scores (~0.09–0.18 range with k=10) so that an exact
975/// name match always ranks above any partial or semantic match.
976const EXACT_MATCH_BOOST: f64 = 0.5;
977
978/// Fuse text + vector hits with Reciprocal Rank Fusion (k=10).
979///
980/// Entity search stays local because it uses k=10 plus exact-match boosting.
981/// Hits in both lists get RRF scores summed. If `query_text` exactly matches
982/// (case-insensitive) an entity's title from the text hits, a bonus of
983/// `EXACT_MATCH_BOOST` is added to ensure exact-name matches dominate.
984/// Sort by fused score, take top-`limit`.
985fn rrf_fuse(
986    text_hits: Vec<TextSearchHit>,
987    vector_hits: Vec<VectorSearchHit>,
988    limit: usize,
989    query_text: &str,
990) -> Vec<SearchHit> {
991    #[derive(Default)]
992    struct Bucket {
993        score: DeterministicScore,
994        source: Option<SearchSource>,
995        title: Option<String>,
996        snippet: Option<String>,
997    }
998
999    let mut buckets: HashMap<Uuid, Bucket> = HashMap::new();
1000
1001    let query_lower = query_text.to_lowercase();
1002    for (i, hit) in text_hits.into_iter().enumerate() {
1003        let rank = i + 1; // RRF is 1-indexed
1004        let entry = buckets.entry(hit.subject_id).or_default();
1005        entry.score = entry.score + rrf_score(rank, RRF_K);
1006        entry.source = Some(match entry.source {
1007            Some(SearchSource::Vector) => SearchSource::Both,
1008            _ => SearchSource::Text,
1009        });
1010        if entry.title.is_none() {
1011            // Apply exact-match boost before storing the title so we only check once.
1012            if let Some(ref title) = hit.title {
1013                if title.to_lowercase() == query_lower {
1014                    entry.score = entry.score + DeterministicScore::from_f64(EXACT_MATCH_BOOST);
1015                }
1016            }
1017            entry.title = hit.title;
1018        }
1019        if entry.snippet.is_none() {
1020            entry.snippet = hit.snippet;
1021        }
1022    }
1023
1024    for (i, hit) in vector_hits.into_iter().enumerate() {
1025        let rank = i + 1;
1026        let entry = buckets.entry(hit.subject_id).or_default();
1027        entry.score = entry.score + rrf_score(rank, RRF_K);
1028        entry.source = Some(match entry.source {
1029            Some(SearchSource::Text) => SearchSource::Both,
1030            _ => SearchSource::Vector,
1031        });
1032    }
1033
1034    let mut hits: Vec<SearchHit> = buckets
1035        .into_iter()
1036        .map(|(id, b)| SearchHit {
1037            entity_id: id,
1038            score: b.score,
1039            source: b.source.expect("each bucket gets a source"),
1040            title: b.title,
1041            snippet: b.snippet,
1042        })
1043        .collect();
1044
1045    hits.sort_by(|a, b| b.score.cmp(&a.score).then(a.entity_id.cmp(&b.entity_id)));
1046    hits.truncate(limit);
1047    hits
1048}
1049
1050#[cfg(test)]
1051mod tests {
1052    use super::*;
1053    use crate::runtime::{KhiveRuntime, NamespaceToken, RuntimeConfig};
1054    use khive_storage::types::{TextSearchHit, VectorSearchHit};
1055    use khive_types::namespace::Namespace;
1056    use lattice_embed::EmbeddingModel;
1057
1058    fn text_hit(id: Uuid, rank: u32, title: &str) -> TextSearchHit {
1059        TextSearchHit {
1060            subject_id: id,
1061            score: DeterministicScore::from_f64(1.0),
1062            rank,
1063            title: Some(title.to_string()),
1064            snippet: Some("...".to_string()),
1065        }
1066    }
1067
1068    fn vector_hit(id: Uuid, rank: u32) -> VectorSearchHit {
1069        VectorSearchHit {
1070            subject_id: id,
1071            score: DeterministicScore::from_f64(0.9),
1072            rank,
1073        }
1074    }
1075
1076    #[test]
1077    fn rrf_fuse_text_only() {
1078        let a = Uuid::new_v4();
1079        let b = Uuid::new_v4();
1080        let text = vec![text_hit(a, 1, "A"), text_hit(b, 2, "B")];
1081        let hits = rrf_fuse(text, vec![], 10, "query");
1082        assert_eq!(hits.len(), 2);
1083        assert_eq!(hits[0].entity_id, a);
1084        assert_eq!(hits[0].source, SearchSource::Text);
1085        assert_eq!(hits[0].title.as_deref(), Some("A"));
1086    }
1087
1088    #[test]
1089    fn rrf_fuse_vector_only() {
1090        let a = Uuid::new_v4();
1091        let hits = rrf_fuse(vec![], vec![vector_hit(a, 1)], 10, "query");
1092        assert_eq!(hits.len(), 1);
1093        assert_eq!(hits[0].source, SearchSource::Vector);
1094        assert!(hits[0].title.is_none());
1095    }
1096
1097    #[test]
1098    fn rrf_fuse_marks_both_when_in_both_lists() {
1099        let id = Uuid::new_v4();
1100        let text = vec![text_hit(id, 1, "A")];
1101        let vec = vec![vector_hit(id, 1)];
1102        let hits = rrf_fuse(text, vec, 10, "query");
1103        assert_eq!(hits.len(), 1);
1104        assert_eq!(hits[0].source, SearchSource::Both);
1105    }
1106
1107    #[test]
1108    fn rrf_fuse_respects_limit() {
1109        let hits: Vec<TextSearchHit> = (0..20)
1110            .map(|i| text_hit(Uuid::new_v4(), i + 1, "x"))
1111            .collect();
1112        let fused = rrf_fuse(hits, vec![], 5, "query");
1113        assert_eq!(fused.len(), 5);
1114    }
1115
1116    #[test]
1117    fn rrf_fuse_orders_higher_score_first() {
1118        // Same UUID in both lists at rank 1 → score 2/(10+1). Different UUIDs → 1/(10+1) each.
1119        let a = Uuid::new_v4();
1120        let b = Uuid::new_v4();
1121        let text = vec![text_hit(a, 1, "A")];
1122        let vec = vec![vector_hit(a, 1), vector_hit(b, 2)];
1123        let hits = rrf_fuse(text, vec, 10, "query");
1124        assert_eq!(hits[0].entity_id, a);
1125        assert_eq!(hits[0].source, SearchSource::Both);
1126        assert!(hits[0].score > hits[1].score);
1127    }
1128
1129    #[test]
1130    fn rrf_fuse_k10_score_spread_exceeds_threshold() {
1131        // With k=10: rank 1 → 1/11 ≈ 0.0909, rank 10 → 1/20 = 0.0500.
1132        // Spread ≈ 0.041, well above the 0.03 minimum required for reliable dedup.
1133        let ids: Vec<Uuid> = (0..10).map(|_| Uuid::new_v4()).collect();
1134        let text: Vec<TextSearchHit> = ids
1135            .iter()
1136            .enumerate()
1137            .map(|(i, &id)| text_hit(id, (i + 1) as u32, "x"))
1138            .collect();
1139        let hits = rrf_fuse(text, vec![], 10, "query");
1140        assert_eq!(hits.len(), 10);
1141        let top_score = hits[0].score.to_f64();
1142        let bottom_score = hits[9].score.to_f64();
1143        let spread = top_score - bottom_score;
1144        assert!(
1145            spread >= 0.03,
1146            "score spread {spread:.4} between rank 1 and rank 10 must be ≥ 0.03 (was {spread:.4})"
1147        );
1148    }
1149
1150    #[test]
1151    fn rrf_fuse_exact_match_boost_elevates_score() {
1152        // An entity whose title exactly matches the query should receive a score
1153        // significantly above a non-matching entity ranked first by text search.
1154        let exact_id = Uuid::new_v4();
1155        let other_id = Uuid::new_v4();
1156        // other_id ranks 1 in text, exact_id ranks 2 — but exact_id matches query.
1157        let text = vec![
1158            text_hit(other_id, 1, "something else"),
1159            text_hit(exact_id, 2, "FlashAttention"),
1160        ];
1161        let hits = rrf_fuse(text, vec![], 10, "flashattention");
1162        assert_eq!(hits.len(), 2);
1163        assert_eq!(
1164            hits[0].entity_id, exact_id,
1165            "exact match must rank first despite being rank-2 in raw text search"
1166        );
1167    }
1168
1169    // ---- embed_batch tests ----
1170
1171    #[test]
1172    fn embed_batch_unconfigured_on_memory_runtime() {
1173        // KhiveRuntime::memory() has no embedding model — embed_batch returns Unconfigured.
1174        let rt = KhiveRuntime::memory().unwrap();
1175        let result = tokio::runtime::Runtime::new()
1176            .unwrap()
1177            .block_on(rt.embed_batch(&[]));
1178        // Empty slice short-circuits before hitting the model check.
1179        assert!(result.is_ok());
1180        assert!(result.unwrap().is_empty());
1181    }
1182
1183    #[test]
1184    fn embed_batch_empty_input_returns_empty_vec() {
1185        // No model needed — empty slice is handled before the embedder is touched.
1186        let rt = KhiveRuntime::memory().unwrap();
1187        let result = tokio::runtime::Runtime::new()
1188            .unwrap()
1189            .block_on(rt.embed_batch(&[]));
1190        assert_eq!(result.unwrap(), Vec::<Vec<f32>>::new());
1191    }
1192
1193    #[test]
1194    fn embed_batch_no_model_non_empty_returns_unconfigured() {
1195        let rt = KhiveRuntime::memory().unwrap();
1196        let texts = vec!["hello".to_string()];
1197        let result = tokio::runtime::Runtime::new()
1198            .unwrap()
1199            .block_on(rt.embed_batch(&texts));
1200        match result {
1201            Err(crate::RuntimeError::Unconfigured(s)) => assert_eq!(s, "embedding_model"),
1202            Err(other) => panic!("expected Unconfigured, got {:?}", other),
1203            Ok(_) => panic!("expected Err, got Ok"),
1204        }
1205    }
1206
1207    #[test]
1208    #[ignore = "loads ~80 MB model; run with --include-ignored"]
1209    fn embed_batch_count_matches_input() {
1210        let config = RuntimeConfig {
1211            db_path: None,
1212            default_namespace: Namespace::parse("test").unwrap(),
1213            embedding_model: Some(EmbeddingModel::AllMiniLmL6V2),
1214            packs: vec!["kg".to_string()],
1215            ..RuntimeConfig::default()
1216        };
1217        let rt = KhiveRuntime::new(config).unwrap();
1218        let texts: Vec<String> = vec!["foo".to_string(), "bar".to_string(), "baz".to_string()];
1219        let result = tokio::runtime::Runtime::new()
1220            .unwrap()
1221            .block_on(rt.embed_batch(&texts));
1222        let embeddings = result.unwrap();
1223        assert_eq!(embeddings.len(), texts.len());
1224    }
1225
1226    #[test]
1227    fn vector_search_requires_embedding_or_text() {
1228        let rt = KhiveRuntime::memory().unwrap();
1229        let tok = NamespaceToken::local();
1230        let result = tokio::runtime::Runtime::new()
1231            .unwrap()
1232            .block_on(rt.vector_search(&tok, None, None, 10, Some(SubstrateKind::Entity)));
1233        match result {
1234            Err(crate::RuntimeError::InvalidInput(msg)) => {
1235                assert!(msg.contains("query_embedding or query_text"), "msg: {msg}");
1236            }
1237            other => panic!("expected InvalidInput, got {other:?}"),
1238        }
1239    }
1240
1241    #[test]
1242    fn vector_search_text_without_model_returns_unconfigured() {
1243        let rt = KhiveRuntime::memory().unwrap();
1244        let tok = NamespaceToken::local();
1245        let result = tokio::runtime::Runtime::new()
1246            .unwrap()
1247            .block_on(rt.vector_search(
1248                &tok,
1249                None,
1250                Some("attention"),
1251                10,
1252                Some(SubstrateKind::Entity),
1253            ));
1254        match result {
1255            Err(crate::RuntimeError::Unconfigured(s)) => assert_eq!(s, "embedding_model"),
1256            other => panic!("expected Unconfigured, got {other:?}"),
1257        }
1258    }
1259
1260    #[test]
1261    #[ignore = "loads ~80 MB model; run with --include-ignored"]
1262    fn embed_batch_vectors_have_expected_dimensions() {
1263        let model = EmbeddingModel::AllMiniLmL6V2;
1264        let config = RuntimeConfig {
1265            db_path: None,
1266            default_namespace: Namespace::parse("test").unwrap(),
1267            embedding_model: Some(model),
1268            packs: vec!["kg".to_string()],
1269            ..RuntimeConfig::default()
1270        };
1271        let rt = KhiveRuntime::new(config).unwrap();
1272        let texts = vec!["hello world".to_string()];
1273        let result = tokio::runtime::Runtime::new()
1274            .unwrap()
1275            .block_on(rt.embed_batch(&texts));
1276        let embeddings = result.unwrap();
1277        assert_eq!(embeddings[0].len(), model.dimensions());
1278    }
1279
1280    // ---- hybrid_search enrichment (issue #147 / #160) ----
1281
1282    #[tokio::test]
1283    async fn hybrid_search_entity_hit_has_title() {
1284        let rt = KhiveRuntime::memory().unwrap();
1285        let tok = NamespaceToken::local();
1286        rt.create_entity(
1287            &tok,
1288            "concept",
1289            None,
1290            "FlashAttention",
1291            Some("IO-aware exact attention using tiling"),
1292            None,
1293            vec![],
1294        )
1295        .await
1296        .unwrap();
1297
1298        let hits = rt
1299            .hybrid_search(&tok, "FlashAttention", None, 10, None, None, &[], None)
1300            .await
1301            .unwrap();
1302
1303        assert!(!hits.is_empty(), "should find the entity");
1304        let hit = &hits[0];
1305        assert!(hit.title.is_some(), "title must be populated");
1306        assert!(
1307            hit.title.as_deref().unwrap().contains("FlashAttention"),
1308            "title must contain entity name"
1309        );
1310    }
1311
1312    // ---- issue #225 regression: predicate pushdown before truncation ----
1313
1314    /// Regression test for issue #225 (entity branch).
1315    ///
1316    /// Scenario: `limit=1`, tag_filter=["target-tag"]. Two entities are inserted:
1317    ///   - "decoy_alpha_beta_gamma": many query tokens → ranks 1 in FTS (dominates).
1318    ///     Does NOT have "target-tag".
1319    ///   - "alpha_beta_gamma target": fewer query tokens → ranks 2 in FTS.
1320    ///     HAS "target-tag".
1321    ///
1322    /// Without predicate pushdown: `fused.truncate(1)` keeps only the decoy. The
1323    /// tag-matching entity is invisible. The test asserts the matching entity IS
1324    /// returned — this assertion fails on the unfixed code (where tags_any is not
1325    /// passed into `query_entities`) and passes after the fix.
1326    ///
1327    /// Isomorphism: reverting `tags_any: tags_any.to_vec()` in the `EntityFilter`
1328    /// inside `hybrid_search` re-breaks this test (the decoy survives `retain` and
1329    /// occupies the single slot, dropping the target).
1330    #[tokio::test]
1331    async fn hybrid_search_tag_filter_pushed_before_truncation() {
1332        let rt = KhiveRuntime::memory().unwrap();
1333        let tok = NamespaceToken::local();
1334
1335        // Decoy: high-ranking FTS hit (content repeats query words), no target tag.
1336        rt.create_entity(
1337            &tok,
1338            "concept",
1339            None,
1340            "alpha beta gamma decoy alpha beta gamma",
1341            Some("alpha beta gamma decoy description alpha beta gamma"),
1342            None,
1343            vec!["other-tag".to_string()],
1344        )
1345        .await
1346        .unwrap();
1347
1348        // Target: lower-ranking FTS hit, has the tag we filter on.
1349        let target = rt
1350            .create_entity(
1351                &tok,
1352                "concept",
1353                None,
1354                "alpha beta gamma target",
1355                Some("alpha beta gamma target description"),
1356                None,
1357                vec!["target-tag".to_string()],
1358            )
1359            .await
1360            .unwrap();
1361
1362        // With limit=1 and tag_filter, the fix must return the target entity despite
1363        // the decoy ranking higher. Without pushdown, the decoy occupies the single
1364        // slot and the target is silently dropped.
1365        let hits = rt
1366            .hybrid_search(
1367                &tok,
1368                "alpha beta gamma",
1369                None,
1370                1,
1371                None,
1372                None,
1373                &["target-tag".to_string()],
1374                None,
1375            )
1376            .await
1377            .unwrap();
1378
1379        assert_eq!(
1380            hits.len(),
1381            1,
1382            "exactly one hit expected (the tag-matching entity)"
1383        );
1384        assert_eq!(
1385            hits[0].entity_id, target.id,
1386            "the tag-filtered entity must be returned even when ranked below limit in raw fusion"
1387        );
1388    }
1389
1390    /// Regression test for issue #225 (entity branch, properties predicate).
1391    ///
1392    /// Scenario: `limit=1`, properties_filter={{"domain": "target"}}. Two entities:
1393    ///   - decoy: high FTS rank, properties {{"domain": "other"}}.
1394    ///   - target: lower FTS rank, properties {{"domain": "target"}}.
1395    ///
1396    /// Without pushdown: decoy fills the slot, target is dropped. With pushdown:
1397    /// only the target survives the properties filter before truncation.
1398    #[tokio::test]
1399    async fn hybrid_search_props_filter_pushed_before_truncation() {
1400        let rt = KhiveRuntime::memory().unwrap();
1401        let tok = NamespaceToken::local();
1402
1403        rt.create_entity(
1404            &tok,
1405            "concept",
1406            None,
1407            "delta epsilon zeta decoy delta epsilon zeta",
1408            Some("delta epsilon zeta decoy description delta epsilon zeta"),
1409            Some(serde_json::json!({"domain": "other"})),
1410            vec![],
1411        )
1412        .await
1413        .unwrap();
1414
1415        let target = rt
1416            .create_entity(
1417                &tok,
1418                "concept",
1419                None,
1420                "delta epsilon zeta target",
1421                Some("delta epsilon zeta target description"),
1422                Some(serde_json::json!({"domain": "target"})),
1423                vec![],
1424            )
1425            .await
1426            .unwrap();
1427
1428        let filter = serde_json::json!({"domain": "target"});
1429        let hits = rt
1430            .hybrid_search(
1431                &tok,
1432                "delta epsilon zeta",
1433                None,
1434                1,
1435                None,
1436                None,
1437                &[],
1438                Some(&filter),
1439            )
1440            .await
1441            .unwrap();
1442
1443        assert_eq!(hits.len(), 1, "exactly one hit expected (properties match)");
1444        assert_eq!(
1445            hits[0].entity_id, target.id,
1446            "the properties-filtered entity must be returned even when ranked below limit"
1447        );
1448    }
1449
1450    // ---- embed intent tests (issue #93) ----
1451
1452    #[test]
1453    #[ignore = "loads ~80 MB model; run with --include-ignored"]
1454    fn minilm_document_and_query_embed_are_identical_no_prefix_model() {
1455        // MiniLM has no instruction prefixes; document and query paths must
1456        // produce byte-identical vectors so that existing stored vectors remain
1457        // comparable after this change.
1458        let model = EmbeddingModel::AllMiniLmL6V2;
1459        let config = RuntimeConfig {
1460            db_path: None,
1461            default_namespace: Namespace::parse("test").unwrap(),
1462            embedding_model: Some(model),
1463            packs: vec!["kg".to_string()],
1464            ..RuntimeConfig::default()
1465        };
1466        let rt = KhiveRuntime::new(config).unwrap();
1467        let text = "attention is all you need".to_string();
1468        let rt_ref = &rt;
1469        let (doc_emb, query_emb) = tokio::runtime::Runtime::new().unwrap().block_on(async {
1470            let d = rt_ref
1471                .embed_document_with_model(&model.to_string(), &text)
1472                .await
1473                .unwrap();
1474            let q = rt_ref
1475                .embed_query_with_model(&model.to_string(), &text)
1476                .await
1477                .unwrap();
1478            (d, q)
1479        });
1480        assert_eq!(
1481            doc_emb, query_emb,
1482            "MiniLM has no instruction prefix: document and query embeds must be identical"
1483        );
1484    }
1485
1486    #[test]
1487    #[ignore = "loads multilingual-e5-small (~90 MB); run with --include-ignored"]
1488    fn e5_document_and_query_embed_differ_instruction_tuned_model() {
1489        // multilingual-e5 prepends "passage: " for documents and "query: " for
1490        // queries. The same raw text must produce different embeddings when the
1491        // correct prefixes are applied, confirming the asymmetric-retrieval
1492        // capability is now exercised.
1493        let model = EmbeddingModel::MultilingualE5Small;
1494        let config = RuntimeConfig {
1495            db_path: None,
1496            default_namespace: Namespace::parse("test").unwrap(),
1497            embedding_model: Some(model),
1498            packs: vec!["kg".to_string()],
1499            ..RuntimeConfig::default()
1500        };
1501        let rt = KhiveRuntime::new(config).unwrap();
1502        let text = "attention is all you need".to_string();
1503        let rt_ref = &rt;
1504        let (doc_emb, query_emb) = tokio::runtime::Runtime::new().unwrap().block_on(async {
1505            let d = rt_ref
1506                .embed_document_with_model(&model.to_string(), &text)
1507                .await
1508                .unwrap();
1509            let q = rt_ref
1510                .embed_query_with_model(&model.to_string(), &text)
1511                .await
1512                .unwrap();
1513            (d, q)
1514        });
1515        assert_ne!(
1516            doc_emb, query_emb,
1517            "multilingual-e5-small uses asymmetric prefixes: document ('passage: ') \
1518             and query ('query: ') embeds of the same text must differ"
1519        );
1520    }
1521
1522    // ---- M-07 regression: backfill reader error must be propagated, not swallowed ----
1523
1524    use crate::embedder_registry::EmbedderProvider;
1525    use lattice_embed::EmbeddingService;
1526
1527    /// A stub embedder that never actually loads weights — used to satisfy the
1528    /// `registered_embedding_model_names` check inside `backfill_missing_embeddings`
1529    /// without triggering a real model load. The test fault-injects a reader error
1530    /// before any embedding call is made, so `embed()` is never reached.
1531    struct StubEmbedderProvider;
1532
1533    #[async_trait::async_trait]
1534    impl EmbedderProvider for StubEmbedderProvider {
1535        fn name(&self) -> &str {
1536            "stub-model-m07"
1537        }
1538
1539        fn dimensions(&self) -> usize {
1540            4
1541        }
1542
1543        async fn build(&self) -> crate::error::RuntimeResult<std::sync::Arc<dyn EmbeddingService>> {
1544            struct StubSvc;
1545            #[async_trait::async_trait]
1546            impl EmbeddingService for StubSvc {
1547                async fn embed(
1548                    &self,
1549                    _texts: &[String],
1550                    _model: lattice_embed::EmbeddingModel,
1551                ) -> std::result::Result<Vec<Vec<f32>>, lattice_embed::EmbedError> {
1552                    Ok(vec![])
1553                }
1554
1555                fn supports_model(&self, _model: lattice_embed::EmbeddingModel) -> bool {
1556                    true
1557                }
1558
1559                fn name(&self) -> &'static str {
1560                    "stub-svc-m07"
1561                }
1562            }
1563            Ok(std::sync::Arc::new(StubSvc))
1564        }
1565    }
1566
1567    /// Regression test for M-07: `backfill_missing_embeddings` must propagate a
1568    /// reader error rather than treating it as "zero rows to embed" (silent swallow).
1569    ///
1570    /// Before the fix: `Err(_) => vec![]` caused the caller to receive `Ok(0)`,
1571    /// silently skipping all embeddings. After the fix: the error is returned via `?`.
1572    ///
1573    /// The fault injection substitutes a `StorageError::Pool` for the result of
1574    /// `sql.reader().await` (i.e., the error originates AT the reader boundary, not
1575    /// before it), so the test exercises the exact `map_err(RuntimeError::Storage)?`
1576    /// lines that the fix introduced. Reverting those lines to `unwrap_or_default()`
1577    /// would swallow the injected error and cause this test to fail.
1578    #[tokio::test]
1579    async fn backfill_reader_error_is_propagated_not_swallowed() {
1580        let rt = KhiveRuntime::memory().unwrap();
1581        rt.register_embedder(StubEmbedderProvider);
1582        let tok = NamespaceToken::local();
1583
1584        // Arm the fault injection: the next backfill call will substitute a
1585        // StorageError at the sql.reader().await boundary, then reset.
1586        super::arm_backfill_reader_fail();
1587
1588        let result = rt.backfill_missing_embeddings(&tok).await;
1589        assert!(
1590            result.is_err(),
1591            "backfill_missing_embeddings must propagate the reader error (got Ok instead)"
1592        );
1593        let err_msg = result.unwrap_err().to_string();
1594        assert!(
1595            err_msg.contains("injected failure"),
1596            "error must originate from the injected reader failure, got: {err_msg}"
1597        );
1598    }
1599}