altair823-org/kebab
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 45 Wiki Activity

refactor(rename): kb crates → kebab — Cargo packages, folders, Rust modules

Browse Source 
프로젝트 이름 `kb` → `kebab` rename 의 첫 단계.

- workspace `Cargo.toml`: members `crates/kb-*` → `crates/kebab-*`,
  repository URL `altair823/kb` → `altair823/kebab`.
- 18 crate 폴더 rename via `git mv` (history 보존).
- 각 crate `Cargo.toml`: `name = "kb-*"` → `"kebab-*"`, path deps
  `../kb-*` → `../kebab-*`.
- 모든 `.rs`: `kb_<id>` snake-case 모듈 path 18 개 (`kb_core`,
  `kb_config`, `kb_app`, `kb_cli`, `kb_eval`, `kb_search`, `kb_chunk`,
  `kb_normalize`, `kb_source_fs`, `kb_parse_md`, `kb_parse_types`,
  `kb_store_sqlite`, `kb_store_vector`, `kb_embed`, `kb_embed_local`,
  `kb_llm`, `kb_llm_local`, `kb_rag`) → `kebab_<id>` 일괄 sed (단어
  경계 \\b 사용해 영어 문장 안의 "kb" 약어 미오염).

CLI binary 이름 (`[[bin]] name = "kb"`), 환경변수 `KB_*`, XDG paths,
tracing target, 그리고 docs sweep 은 다음 commit 에서.

## 검증

- `cargo check --workspace` clean — 모든 crate 빌드 통과 후 commit.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
altair823
2026-05-02 03:28:08 +00:00
parent 2aecbf3d9f
commit 911fb49550
143 changed files with 727 additions and 727 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
crates/kebab-search/Cargo.toml Normal file
View File

@@ -0,0 +1,35 @@
[package]
name = "kebab-search"
version = { workspace = true }
edition = { workspace = true }
rust-version = { workspace = true }
license = { workspace = true }
repository = { workspace = true }
description = "Retriever implementations for kb (P2-2 lexical FTS5; P3 vector / hybrid will follow)"
[dependencies]
kebab-core = { path = "../kebab-core" }
kebab-config = { path = "../kebab-config" }
kebab-store-sqlite = { path = "../kebab-store-sqlite" }
# P3-4 hybrid retriever wraps a `dyn VectorStore` (typically backed by
# `kb-store-vector::LanceVectorStore`) and a `dyn Embedder` (any P3-2
# adapter). Listed as a runtime dep so callers can construct
# `VectorRetriever::new` against the trait objects without a concrete
# adapter — the concrete adapter (`kb-embed-local`) stays out of this
# crate per the spec's Forbidden deps list.
kebab-store-vector = { path = "../kebab-store-vector" }
kebab-embed = { path = "../kebab-embed" }
rusqlite = { workspace = true }
globset = { workspace = true }
serde_json = { workspace = true }
tracing = { workspace = true }
thiserror = { workspace = true }
anyhow = { workspace = true }
[dev-dependencies]
tempfile = { workspace = true }
# Hybrid integration tests inject a `MockEmbedder` (kb-embed `mock`
# feature) and stand up a real `LanceVectorStore` on a tmp directory.
# The mock-retriever unit tests (the bulk of the hybrid suite) do not
# need either, but the integration / snapshot lane does.
kebab-embed = { path = "../kebab-embed", features = ["mock"] }

74
crates/kebab-search/src/citation_helper.rs Normal file
View File

@@ -0,0 +1,74 @@
//! Shared helpers for building `kebab_core::Citation` values from a
//! chunk's first `SourceSpan`.
//!
//! Both the lexical and vector retrievers join against the same
//! `chunks.source_spans_json` column and need identical mapping logic
//! so cross-mode citation strings round-trip byte-identically (a
//! requirement for the hybrid retriever's tie-break on chunk_id and
//! for the `search --explain` output documented in design §0 Q3 and
//! §1.6). Living here means a future PDF / image / audio extractor can
//! enrich the mapping in one place rather than two.
use kebab_core::{Citation, SourceSpan, WorkspacePath};
/// Build a `Citation` from the chunk's first `SourceSpan`. P1 markdown
/// only emits `Line`, so the other variants are mostly defensive — we
/// forward them as faithfully as possible so a future PDF / image
/// extractor can flow through without churn.
///
/// `chunk_id` is taken only for diagnostic logging when the span shape
/// has no Citation mapping (`Byte`-spans, empty arrays).
pub(crate) fn citation_from_first_span(
chunk_id: &str,
path: WorkspacePath,
section: Option<String>,
first_span: Option<&SourceSpan>,
) -> Citation {
match first_span {
Some(SourceSpan::Line { start, end }) => Citation::Line {
path,
start: *start,
end: *end,
section,
},
Some(SourceSpan::Page { page, .. }) => Citation::Page {
path,
page: *page,
section,
},
Some(SourceSpan::Region { x, y, w, h }) => Citation::Region {
path,
x: *x,
y: *y,
w: *w,
h: *h,
},
Some(SourceSpan::Time { start_ms, end_ms }) => Citation::Time {
path,
start_ms: *start_ms,
end_ms: *end_ms,
speaker: None,
},
// Byte-spans don't have a Citation variant. Fall back to a
// Line citation pointing at the document head — better than
// fabricating a position. Spans-empty falls into the same
// branch.
other @ (Some(SourceSpan::Byte { .. }) | None) => {
let span_shape = match other {
Some(_) => "Byte",
None => "empty array",
};
tracing::warn!(
chunk_id,
span_shape,
"kb-search: SourceSpan has no Citation mapping; falling back to Line {{1, 1}}"
);
Citation::Line {
path,
start: 1,
end: 1,
section,
}
}
}
}

632
crates/kebab-search/src/hybrid.rs Normal file
View File

@@ -0,0 +1,632 @@
//! Hybrid retriever — design §3.7 / §6.4 / §0 Q3 / §1.6.
//!
//! Composes a lexical and a vector retriever (both `dyn Retriever`)
//! and dispatches by `SearchMode`. For `Hybrid`, results are fused via
//! Reciprocal Rank Fusion (RRF):
//!
//! ```text
//! score(c) = Σ_{m ∈ {lex, vec}} 1 / (k_rrf + rank_m(c))
//! ```
//!
//! where `rank_m(c)` is the 1-based rank of chunk `c` in retriever
//! `m`'s output (chunks not appearing in `m` contribute 0).
//!
//! Each `SearchHit.retrieval` is rebuilt with the per-mode scores /
//! ranks the fusion observed, so `kb search --explain` (§1.6) can
//! show users exactly which retriever contributed what to the final
//! ordering.
use std::collections::HashMap;
use std::sync::Arc;
use anyhow::Result;
use kebab_core::{
IndexVersion, RetrievalDetail, Retriever, SearchHit, SearchMode, SearchQuery,
};
/// Default `k_rrf` if `kb-config::SearchCfg::rrf_k` is misconfigured.
/// Matches §6.4's documented default (60).
const DEFAULT_K_RRF: u32 = 60;
/// When fanning out for hybrid fusion we ask each side for `k *
/// HYBRID_FANOUT_MULTIPLIER` candidates so the disjoint set of
/// chunks (those a single retriever surfaces but the other does not)
/// is wide enough to feed a useful fused top-k.
///
/// `2` is the spec-suggested floor; raising it helps recall on
/// adversarial corpora at linear cost. Documented in
/// `tasks/p3/p3-4-hybrid-fusion.md` "Risks / notes".
const HYBRID_FANOUT_MULTIPLIER: usize = 2;
/// Default `k` when `SearchQuery::k == 0`. Mirrors §6.4 default_k=10.
const DEFAULT_K: usize = 10;
/// Fusion algorithm. Today only Reciprocal Rank Fusion is supported;
/// listing as an enum so future score-calibration policies (P+) can
/// land without an API break.
#[derive(Clone, Copy, Debug)]
pub enum FusionPolicy {
/// Reciprocal Rank Fusion. `k_rrf` is the standard rank-bias
/// hyperparameter (§6.4); larger values flatten the rank-bias
/// curve, smaller values privilege top-of-list hits.
Rrf { k_rrf: u32 },
}
/// Hybrid retriever composing a lexical and a vector retriever.
///
/// For chunks that appear in both retrievers, the lexical-side hit
/// supplies `snippet`, `citation`, `heading_path`, `chunker_version`,
/// and `embedding_model` — lexical search has FTS5 highlighting that's
/// more user-relevant than the vector retriever's truncated text.
/// Vector-only chunks fall through to the vector hit's data verbatim.
/// This matches `kb search --explain` (§1.6) expectations for snippet
/// provenance.
pub struct HybridRetriever {
lexical: Arc<dyn Retriever>,
vector: Arc<dyn Retriever>,
fusion: FusionPolicy,
/// Default `k` for queries that arrive with `k == 0`. Pulled from
/// `config.search.default_k` at construction.
default_k: usize,
}
impl HybridRetriever {
/// Construct from a `kb-config` Config + the two underlying
/// retrievers. Reads `config.search.hybrid_fusion` (only `"rrf"`
/// is recognised today) and `config.search.rrf_k`.
pub fn new(
config: &kebab_config::Config,
lexical: Arc<dyn Retriever>,
vector: Arc<dyn Retriever>,
) -> Self {
let fusion = parse_fusion(&config.search.hybrid_fusion, config.search.rrf_k);
let default_k = if config.search.default_k == 0 {
DEFAULT_K
} else {
config.search.default_k
};
// Surface mismatched index_version up front so users see it
// (e.g. lexical at v2, vector at v1 means a stale index that
// the user should refresh). Spec line 144 calls this out as
// a "flag at construction".
let lex_iv = lexical.index_version();
let vec_iv = vector.index_version();
if lex_iv.0 != vec_iv.0 {
tracing::warn!(
target: "kb-search",
lexical_index = %lex_iv.0,
vector_index = %vec_iv.0,
"kb-search hybrid: lexical and vector index_version differ; consider re-indexing"
);
}
Self {
lexical,
vector,
fusion,
default_k,
}
}
/// Construct with explicit policy / `k`. Used by tests that want
/// to pin RRF parameters without going through `kb-config`.
pub fn with_policy(
lexical: Arc<dyn Retriever>,
vector: Arc<dyn Retriever>,
fusion: FusionPolicy,
default_k: usize,
) -> Self {
Self {
lexical,
vector,
fusion,
default_k: if default_k == 0 { DEFAULT_K } else { default_k },
}
}
}
impl Retriever for HybridRetriever {
fn search(&self, query: &SearchQuery) -> Result<Vec<SearchHit>> {
match query.mode {
SearchMode::Lexical => self.lexical.search(query),
SearchMode::Vector => self.vector.search(query),
SearchMode::Hybrid => self.fuse(query),
}
}
fn index_version(&self) -> IndexVersion {
// Composite token so callers (e.g. snapshot tests) can detect
// either side drifting without inspecting both retrievers.
let lex = self.lexical.index_version().0;
let vec = self.vector.index_version().0;
IndexVersion(format!("hybrid:{lex}+{vec}"))
}
}
impl HybridRetriever {
fn fuse(&self, query: &SearchQuery) -> Result<Vec<SearchHit>> {
let target_k = if query.k == 0 { self.default_k } else { query.k };
// Fanout: ask each retriever for `target_k * MULTIPLIER` so
// the disjoint set of candidates is wide enough. The two
// per-side queries are identical (same text, k, mode, filters);
// only the dispatch differs, so we share one `SearchQuery`.
let fanout_k = target_k.saturating_mul(HYBRID_FANOUT_MULTIPLIER);
let lex_query = SearchQuery {
k: fanout_k,
..query.clone()
};
let lex_hits = self.lexical.search(&lex_query)?;
let vec_hits = self.vector.search(&lex_query)?;
tracing::debug!(
lex = lex_hits.len(),
vec = vec_hits.len(),
target_k,
"kb-search hybrid: pre-fusion candidate counts"
);
// Build (chunk_id → (rank, hit)) maps. The rank stored here
// is the `rank` field on each retriever's output, which is
// already 1-based by both LexicalRetriever and VectorRetriever
// (and any well-behaved Retriever should mirror).
let lex_index: HashMap<String, (u32, SearchHit)> = lex_hits
.into_iter()
.map(|h| (h.chunk_id.0.clone(), (h.rank, h)))
.collect();
let vec_index: HashMap<String, (u32, SearchHit)> = vec_hits
.into_iter()
.map(|h| (h.chunk_id.0.clone(), (h.rank, h)))
.collect();
// Union of chunk_ids from both sides.
let mut all_ids: Vec<String> = Vec::with_capacity(lex_index.len() + vec_index.len());
for k in lex_index.keys() {
all_ids.push(k.clone());
}
for k in vec_index.keys() {
if !lex_index.contains_key(k) {
all_ids.push(k.clone());
}
}
// Compute fused score per chunk.
//
// Raw RRF: `Σ 1/(k_rrf + rank_m(c))` over the retrievers a chunk
// appears in. With two retrievers the raw upper bound is
// `2/(k_rrf + 1)` — at k_rrf=60 that's only ≈0.0328, which makes
// a single `config.rag.score_gate` default of 0.05 silently
// refuse every hybrid query (and is incomparable with lexical /
// vector `fusion_score` already in [0, 1]).
//
// Normalize by the theoretical max so `fusion_score` lives in
// [0, 1] across all three SearchModes. The normalization factor
// is `num_retrievers / (k_rrf + 1)`. With both retrievers
// contributing rank=1 the normalized score is exactly 1.0;
// chunks present in only one retriever cap at ≈0.5 (≈ 1 / 2);
// all other rank combinations fall in between. RRF's rank-
// ordering invariants are preserved (we divide every score by
// the same positive constant), so the sort + tiebreak path is
// unchanged. Wire schema label `fusion_score` keeps its slot in
// `RetrievalDetail`; only the magnitude shifts.
let FusionPolicy::Rrf { k_rrf } = self.fusion;
let k_rrf_f = f64::from(k_rrf);
// Both retrievers can contribute, so the per-mode RRF max is
// 2 / (k_rrf + 1). Even when a chunk lands in only one mode, we
// still divide by this same constant — the score then caps
// around 0.5 which is exactly the "half-aligned" semantic we
// want users to compare against `score_gate`.
let rrf_normalizer = 2.0_f64 / (k_rrf_f + 1.0);
struct Scored {
chunk_id: String,
rrf: f64,
lex_rank: Option<u32>,
vec_rank: Option<u32>,
}
let mut scored: Vec<Scored> = all_ids
.into_iter()
.map(|cid| {
let lex_rank = lex_index.get(&cid).map(|(r, _)| *r);
let vec_rank = vec_index.get(&cid).map(|(r, _)| *r);
let mut rrf = 0.0_f64;
if let Some(r) = lex_rank {
rrf += 1.0 / (k_rrf_f + f64::from(r));
}
if let Some(r) = vec_rank {
rrf += 1.0 / (k_rrf_f + f64::from(r));
}
rrf /= rrf_normalizer;
Scored {
chunk_id: cid,
rrf,
lex_rank,
vec_rank,
}
})
.collect();
// Sort: rrf DESC, then lex_rank ASC (None last), then chunk_id ASC.
// f64 ordering uses `total_cmp` so NaN stays deterministic
// (won't occur today — k_rrf > 0 → denominators > 0 — but
// total_cmp keeps the sort stable under future tweaks).
scored.sort_by(|a, b| {
b.rrf
.total_cmp(&a.rrf)
.then_with(|| {
let am = a.lex_rank.unwrap_or(u32::MAX);
let bm = b.lex_rank.unwrap_or(u32::MAX);
am.cmp(&bm)
})
.then_with(|| a.chunk_id.cmp(&b.chunk_id))
});
// Build final SearchHits, taking the top `target_k`.
let mut hits: Vec<SearchHit> = Vec::with_capacity(target_k.min(scored.len()));
let mut rank: u32 = 0;
for s in scored.into_iter().take(target_k) {
// Pull the underlying hit. Prefer the lexical side when
// available — its snippet has FTS5 highlighting which
// gives users the most useful preview. Fall back to
// vector if the chunk only appeared in vector results.
let mut base = match (lex_index.get(&s.chunk_id), vec_index.get(&s.chunk_id)) {
(Some((_, lex)), _) => lex.clone(),
(None, Some((_, vec))) => vec.clone(),
// `all_ids` is the union of `lex_index` and
// `vec_index` keys, so this arm cannot fire.
(None, None) => {
unreachable!("chunk_id was in union but absent from both indices")
}
};
// `unwrap_or(fusion_score)` covers a defensive-coding case
// that doesn't arise today: when a chunk only appears in
// one retriever, RRF sums a single term so `fusion_score`
// already equals that side's normalized score, making the
// fallback harmless.
let lex_score = lex_index
.get(&s.chunk_id)
.map(|(_, h)| h.retrieval.lexical_score.unwrap_or(h.retrieval.fusion_score));
let vec_score = vec_index
.get(&s.chunk_id)
.map(|(_, h)| h.retrieval.vector_score.unwrap_or(h.retrieval.fusion_score));
rank = rank.saturating_add(1);
base.rank = rank;
base.retrieval = RetrievalDetail {
method: SearchMode::Hybrid,
// RRF is computed in f64 inside `fuse` and cast to f32
// here at the boundary. `1/(k_rrf+rank)` is bounded
// roughly in `(0, 2/k_rrf]` (≤ ~0.033 at k_rrf=60), so
// the magnitude is well within f32 range and f32
// precision is more than sufficient for ranking.
fusion_score: s.rrf as f32,
lexical_score: lex_score,
vector_score: vec_score,
lexical_rank: s.lex_rank,
vector_rank: s.vec_rank,
};
hits.push(base);
}
tracing::debug!(rows = hits.len(), "kb-search hybrid: search done");
Ok(hits)
}
}
/// Parse the `hybrid_fusion` config string into a [`FusionPolicy`].
/// Today only `"rrf"` is recognised; anything else falls back to RRF
/// with a warn log so misconfiguration is visible but not fatal.
fn parse_fusion(name: &str, k_rrf: u32) -> FusionPolicy {
let k = if k_rrf == 0 { DEFAULT_K_RRF } else { k_rrf };
match name {
"rrf" => FusionPolicy::Rrf { k_rrf: k },
other => {
tracing::warn!(
target: "kb-search",
policy = other,
"kb-search hybrid: unknown fusion policy; falling back to RRF"
);
FusionPolicy::Rrf { k_rrf: k }
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use kebab_core::{
ChunkId, ChunkerVersion, Citation, DocumentId, IndexVersion, SearchFilters,
SearchHit, SearchMode, WorkspacePath,
};
use std::sync::Mutex;
/// Test double: returns a canned `Vec<SearchHit>` and records
/// every call so we can assert delegation.
struct CannedRetriever {
hits: Vec<SearchHit>,
calls: Mutex<Vec<SearchQuery>>,
version: IndexVersion,
}
impl CannedRetriever {
fn new(hits: Vec<SearchHit>, version: &str) -> Self {
Self {
hits,
calls: Mutex::new(Vec::new()),
version: IndexVersion(version.to_string()),
}
}
}
impl Retriever for CannedRetriever {
fn search(&self, query: &SearchQuery) -> Result<Vec<SearchHit>> {
self.calls.lock().unwrap().push(query.clone());
Ok(self.hits.clone())
}
fn index_version(&self) -> IndexVersion {
self.version.clone()
}
}
fn wp(p: &str) -> WorkspacePath {
WorkspacePath::new(p.to_string()).unwrap()
}
/// Build a synthetic `SearchHit`. Most fields take inert defaults
/// because the hybrid logic only reads `chunk_id`, `rank`,
/// `retrieval.{lexical,vector}_score`, and (transitively) the rest
/// when building the fused output.
fn mk_hit(
chunk_id: &str,
rank: u32,
method: SearchMode,
score: f32,
) -> SearchHit {
let cid = ChunkId(chunk_id.to_string());
let did = DocumentId(format!("d-{chunk_id}"));
let path = wp(&format!("notes/{chunk_id}.md"));
SearchHit {
rank,
chunk_id: cid,
doc_id: did,
doc_path: path.clone(),
heading_path: vec![],
section_label: None,
snippet: format!("snippet for {chunk_id}"),
citation: Citation::Line {
path,
start: 1,
end: 1,
section: None,
},
retrieval: RetrievalDetail {
method,
fusion_score: score,
lexical_score: matches!(method, SearchMode::Lexical | SearchMode::Hybrid)
.then_some(score),
vector_score: matches!(method, SearchMode::Vector | SearchMode::Hybrid)
.then_some(score),
lexical_rank: matches!(method, SearchMode::Lexical | SearchMode::Hybrid)
.then_some(rank),
vector_rank: matches!(method, SearchMode::Vector | SearchMode::Hybrid)
.then_some(rank),
},
index_version: IndexVersion("v1".to_string()),
embedding_model: None,
chunker_version: ChunkerVersion("v1".to_string()),
}
}
fn rrf_policy(k_rrf: u32) -> FusionPolicy {
FusionPolicy::Rrf { k_rrf }
}
fn make_query(mode: SearchMode, k: usize) -> SearchQuery {
SearchQuery {
text: "rust".to_string(),
mode,
k,
filters: SearchFilters::default(),
}
}
#[test]
fn hybrid_lexical_mode_delegates_to_lexical() {
let lex_hits = vec![mk_hit("aaaa", 1, SearchMode::Lexical, 0.9)];
let lex = Arc::new(CannedRetriever::new(lex_hits.clone(), "lex-v1"));
let vec = Arc::new(CannedRetriever::new(vec![], "vec-v1"));
let h = HybridRetriever::with_policy(lex.clone(), vec.clone(), rrf_policy(60), 5);
let out = h.search(&make_query(SearchMode::Lexical, 5)).unwrap();
assert_eq!(out, lex_hits, "lexical mode must pass through verbatim");
assert_eq!(lex.calls.lock().unwrap().len(), 1, "lexical called once");
assert_eq!(vec.calls.lock().unwrap().len(), 0, "vector NOT called");
}
#[test]
fn hybrid_vector_mode_delegates_to_vector() {
let vec_hits = vec![mk_hit("bbbb", 1, SearchMode::Vector, 0.8)];
let lex = Arc::new(CannedRetriever::new(vec![], "lex-v1"));
let vec = Arc::new(CannedRetriever::new(vec_hits.clone(), "vec-v1"));
let h = HybridRetriever::with_policy(lex.clone(), vec.clone(), rrf_policy(60), 5);
let out = h.search(&make_query(SearchMode::Vector, 5)).unwrap();
assert_eq!(out, vec_hits, "vector mode must pass through verbatim");
assert_eq!(lex.calls.lock().unwrap().len(), 0, "lexical NOT called");
assert_eq!(vec.calls.lock().unwrap().len(), 1, "vector called once");
}
#[test]
fn hybrid_chunk_only_in_lexical_keeps_vector_none() {
// Chunk X is in lexical only.
let lex = Arc::new(CannedRetriever::new(
vec![mk_hit("xxxx", 1, SearchMode::Lexical, 0.9)],
"lex-v1",
));
let vec = Arc::new(CannedRetriever::new(
vec![mk_hit("yyyy", 1, SearchMode::Vector, 0.8)],
"vec-v1",
));
let h = HybridRetriever::with_policy(lex, vec, rrf_policy(60), 5);
let out = h.search(&make_query(SearchMode::Hybrid, 5)).unwrap();
// Both X and Y are present.
let xx = out.iter().find(|h| h.chunk_id.0 == "xxxx").unwrap();
assert_eq!(xx.retrieval.method, SearchMode::Hybrid);
assert!(xx.retrieval.lexical_score.is_some());
assert_eq!(xx.retrieval.vector_score, None);
assert_eq!(xx.retrieval.lexical_rank, Some(1));
assert_eq!(xx.retrieval.vector_rank, None);
assert!(xx.retrieval.fusion_score > 0.0);
let yy = out.iter().find(|h| h.chunk_id.0 == "yyyy").unwrap();
assert_eq!(yy.retrieval.lexical_score, None);
assert!(yy.retrieval.vector_score.is_some());
assert_eq!(yy.retrieval.lexical_rank, None);
assert_eq!(yy.retrieval.vector_rank, Some(1));
}
#[test]
fn rrf_formula_matches_known_value() {
// chunk A appears at lexical rank 1, vector rank 2; k_rrf=60.
// Raw RRF: 1/(60+1) + 1/(60+2) = 1/61 + 1/62.
// After normalization by `2 / (60 + 1)` (theoretical max with
// both retrievers contributing rank=1), the score lives in
// [0, 1]: `(1/61 + 1/62) / (2/61) = 0.5 + 61/124 ≈ 0.9919`.
let raw = 1.0_f64 / 61.0 + 1.0_f64 / 62.0;
let expected = raw / (2.0_f64 / 61.0);
let lex = Arc::new(CannedRetriever::new(
vec![mk_hit("aaaa", 1, SearchMode::Lexical, 0.5)],
"lex-v1",
));
let vec_hits = vec![
mk_hit("zzzz", 1, SearchMode::Vector, 0.9),
mk_hit("aaaa", 2, SearchMode::Vector, 0.7),
];
let vec = Arc::new(CannedRetriever::new(vec_hits, "vec-v1"));
let h = HybridRetriever::with_policy(lex, vec, rrf_policy(60), 5);
let out = h.search(&make_query(SearchMode::Hybrid, 5)).unwrap();
let a = out.iter().find(|h| h.chunk_id.0 == "aaaa").unwrap();
let actual = a.retrieval.fusion_score as f64;
// Tolerance: the score is computed in f64 and cast to f32 at
// the API boundary, so any discrepancy must fit within f32
// precision. `1e-7` is below `f32::EPSILON` (~1.19e-7), which
// makes the check brittle on edge cases. Use a small multiple
// of EPSILON to stay robust.
let tol = f64::from(f32::EPSILON) * 10.0;
assert!(
(actual - expected).abs() < tol,
"RRF score {actual} drifted from expected {expected} (tol {tol})"
);
}
#[test]
fn hybrid_tiebreak_prefers_lower_lexical_rank_then_chunk_id() {
// Construct two chunks with identical fused scores.
// Strategy: A appears at lex rank 2 only → score = 1/62.
// B appears at vec rank 2 only → score = 1/62.
// Tie-break: lex_rank ascending (Some(2) < None), so A wins.
let lex = Arc::new(CannedRetriever::new(
vec![
mk_hit("zzzz", 1, SearchMode::Lexical, 0.9), // rank 1: high RRF, leader
mk_hit("aaaa", 2, SearchMode::Lexical, 0.5), // rank 2
],
"lex-v1",
));
let vec = Arc::new(CannedRetriever::new(
vec![
mk_hit("zzzz", 1, SearchMode::Vector, 0.9),
mk_hit("bbbb", 2, SearchMode::Vector, 0.5),
],
"vec-v1",
));
let h = HybridRetriever::with_policy(lex, vec, rrf_policy(60), 5);
let out = h.search(&make_query(SearchMode::Hybrid, 5)).unwrap();
// zzzz has both ranks → strictly higher RRF → rank 1.
assert_eq!(out[0].chunk_id.0, "zzzz");
// aaaa and bbbb both have a single rank-2 contribution → identical
// RRF. Tie-break: aaaa has lex_rank=Some(2), bbbb has lex_rank=None,
// so aaaa comes first.
assert_eq!(out[1].chunk_id.0, "aaaa");
assert_eq!(out[2].chunk_id.0, "bbbb");
// Now construct two chunks with identical lex rank to verify
// the chunk_id tie-break. CannedRetriever can't produce two
// hits at the same rank via mk_hit's normal flow, so we patch
// `retrieval.lexical_rank` directly after construction.
let mut tied_a = mk_hit("aaaa", 2, SearchMode::Lexical, 0.4);
tied_a.retrieval.lexical_rank = Some(2);
let mut tied_b = mk_hit("bbbb", 2, SearchMode::Lexical, 0.4);
tied_b.retrieval.lexical_rank = Some(2);
let lex3 = Arc::new(CannedRetriever::new(
vec![tied_a, tied_b],
"lex-v1",
));
let vec3 = Arc::new(CannedRetriever::new(vec![], "vec-v1"));
let h3 = HybridRetriever::with_policy(lex3, vec3, rrf_policy(60), 5);
let out3 = h3.search(&make_query(SearchMode::Hybrid, 5)).unwrap();
// Same lex_rank=2 → tie-break on chunk_id ascending: aaaa < bbbb.
assert_eq!(out3[0].chunk_id.0, "aaaa");
assert_eq!(out3[1].chunk_id.0, "bbbb");
}
#[test]
fn hybrid_index_version_is_composite() {
let lex = Arc::new(CannedRetriever::new(vec![], "lex-v1"));
let vec = Arc::new(CannedRetriever::new(vec![], "vec-v2"));
let h = HybridRetriever::with_policy(lex, vec, rrf_policy(60), 5);
assert_eq!(h.index_version().0, "hybrid:lex-v1+vec-v2");
}
#[test]
fn hybrid_disjoint_recall_returns_all_when_k_large_enough() {
// lex returns [A, B], vec returns [C, D]; k=4 → all 4 in result.
let lex = Arc::new(CannedRetriever::new(
vec![
mk_hit("aaaa", 1, SearchMode::Lexical, 0.9),
mk_hit("bbbb", 2, SearchMode::Lexical, 0.7),
],
"lex-v1",
));
let vec = Arc::new(CannedRetriever::new(
vec![
mk_hit("cccc", 1, SearchMode::Vector, 0.9),
mk_hit("dddd", 2, SearchMode::Vector, 0.7),
],
"vec-v1",
));
let h = HybridRetriever::with_policy(lex, vec, rrf_policy(60), 4);
let out = h.search(&make_query(SearchMode::Hybrid, 4)).unwrap();
let mut ids: Vec<&str> = out.iter().map(|h| h.chunk_id.0.as_str()).collect();
ids.sort();
assert_eq!(ids, vec!["aaaa", "bbbb", "cccc", "dddd"]);
}
#[test]
fn hybrid_zero_k_uses_default() {
// With query.k=0, hybrid should use the configured default_k.
let lex = Arc::new(CannedRetriever::new(
(0..20)
.map(|i| mk_hit(&format!("c{i:04}"), i + 1, SearchMode::Lexical, 0.5))
.collect(),
"lex-v1",
));
let vec = Arc::new(CannedRetriever::new(vec![], "vec-v1"));
let h = HybridRetriever::with_policy(lex, vec, rrf_policy(60), 7);
let out = h.search(&make_query(SearchMode::Hybrid, 0)).unwrap();
assert_eq!(out.len(), 7);
}
#[test]
fn parse_fusion_falls_back_to_rrf_on_unknown() {
let p = parse_fusion("nonsense", 60);
let FusionPolicy::Rrf { k_rrf } = p;
assert_eq!(k_rrf, 60);
}
#[test]
fn parse_fusion_zero_k_falls_back_to_default() {
let FusionPolicy::Rrf { k_rrf } = parse_fusion("rrf", 0);
assert_eq!(k_rrf, DEFAULT_K_RRF);
}
}

596
crates/kebab-search/src/lexical.rs Normal file
View File

@@ -0,0 +1,596 @@
//! Lexical (FTS5 + bm25) retriever — design §3.7 / §1.5 / §2.2 / §6.4.
//!
//! Owns the SQL pattern documented in `tasks/p2/p2-2-lexical-retriever.md`
//! and constructs `kebab_core::SearchHit` values directly from the joined
//! `chunks_fts` / `chunks` / `documents` rows. Reads only — never mutates
//! the underlying SQLite file.
use std::sync::Arc;
use anyhow::{Context, Result};
use globset::GlobMatcher;
use kebab_core::{
ChunkId, ChunkerVersion, DocumentId, IndexVersion, RetrievalDetail, Retriever,
SearchFilters, SearchHit, SearchMode, SearchQuery, SourceSpan, TrustLevel,
WorkspacePath,
};
use kebab_store_sqlite::SqliteStore;
use rusqlite::{params_from_iter, Connection, Row, ToSql};
use crate::citation_helper::citation_from_first_span;
// ── Tunables ─────────────────────────────────────────────────────────────
/// FTS5 hard limit on the `snippet()` `nToken` argument.
/// See SQLite's FTS5 docs: snippet() rejects nToken > 64.
const FTS5_SNIPPET_MAX_WORDS: usize = 64;
/// Floor for the snippet word budget. `snippet_chars / 4` may yield 0 for
/// pathologically small configs; we always ask FTS5 for at least one word
/// so it can still return something matchable for the test harness.
const FTS5_SNIPPET_MIN_WORDS: usize = 1;
/// Default `k` when `SearchQuery::k == 0`. Mirrors §6.4 default_k=10.
const DEFAULT_K: usize = 10;
/// When `path_glob` is set we have to over-fetch and post-filter in Rust,
/// because SQLite's GLOB operator treats `*` as "any chars including `/`",
/// which contradicts the design rule that `*` must NOT cross path
/// separators. Empirically `+128` is generous for any realistic workspace
/// and bounded enough to keep memory predictable.
const PATH_GLOB_OVERFETCH: usize = 128;
// ── Public surface ───────────────────────────────────────────────────────
/// Lexical retriever backed by SQLite FTS5 + bm25.
pub struct LexicalRetriever {
store: Arc<SqliteStore>,
index_version: IndexVersion,
/// Number of `snippet()` words derived from `kb-config::search.snippet_chars`,
/// clamped into `[FTS5_SNIPPET_MIN_WORDS, FTS5_SNIPPET_MAX_WORDS]`.
snippet_words: usize,
/// Hard cap on the returned snippet's character length per design §6.4.
snippet_chars: usize,
}
impl LexicalRetriever {
/// Construct with default settings derived from `kb-config`'s defaults.
/// Snippet width is computed from `Config::defaults().search.snippet_chars`.
pub fn new(store: Arc<SqliteStore>, index_version: IndexVersion) -> Self {
let cfg = kebab_config::Config::defaults();
Self::with_settings(store, index_version, cfg.search.snippet_chars)
}
/// Construct with explicit `snippet_chars`. Used by tests / callers
/// that have already loaded a `Config`.
pub fn with_settings(
store: Arc<SqliteStore>,
index_version: IndexVersion,
snippet_chars: usize,
) -> Self {
// Heuristic: 1 token ≈ 4 chars (English-leaning estimate; Korean
// tokens average shorter, so the cap-by-chars trim below is what
// actually enforces the contract). The `/4` keeps us well below
// FTS5's nToken=64 limit for typical snippet_chars=220 budgets.
let raw = snippet_chars / 4;
let snippet_words = raw.clamp(FTS5_SNIPPET_MIN_WORDS, FTS5_SNIPPET_MAX_WORDS);
Self {
store,
index_version,
snippet_words,
snippet_chars,
}
}
}
impl Retriever for LexicalRetriever {
fn search(&self, query: &SearchQuery) -> Result<Vec<SearchHit>> {
let match_opt = build_match_string(&query.text);
let k = if query.k == 0 { DEFAULT_K } else { query.k };
let filters = &query.filters;
// One-line summary at request entry. Filter shape only — no
// tag/lang/path values, which could be PII-sensitive.
tracing::debug!(
match_str = match_opt.as_deref().unwrap_or("<empty>"),
tags_any = filters.tags_any.len(),
has_lang = filters.lang.is_some(),
has_trust_min = filters.trust_min.is_some(),
has_path_glob = filters.path_glob.is_some(),
k,
"kb-search lexical: search start"
);
// Empty / whitespace-only query → nothing to do. Per spec we
// succeed with an empty hit list rather than erroring.
let match_str = match match_opt {
Some(s) => s,
None => return Ok(Vec::new()),
};
// Pre-compile the path_glob once. The `Glob` produced rejects
// syntactically invalid patterns at construction time so the
// caller gets a clear error rather than a silent empty result.
let path_matcher = match &filters.path_glob {
Some(g) => Some(compile_glob(g)?),
None => None,
};
// Fetch budget: when post-filtering by glob we need to over-fetch
// so that the final `take(k)` still has enough rows after culling.
let fetch_limit = if path_matcher.is_some() {
k.saturating_add(PATH_GLOB_OVERFETCH)
} else {
k
};
let conn = self.store.read_conn();
let raw_rows = run_query(
&conn,
&match_str,
self.snippet_words,
filters,
fetch_limit,
)?;
let mut hits: Vec<SearchHit> = Vec::with_capacity(raw_rows.len().min(k));
let mut rank: u32 = 0;
for row in raw_rows {
// Path glob is the only filter we evaluate in Rust because the
// semantics differ from SQLite's GLOB (no `/` crossing).
if let Some(m) = &path_matcher {
if !m.is_match(&row.workspace_path) {
continue;
}
}
rank = rank.saturating_add(1);
let hit = build_hit(row, rank, &self.index_version, self.snippet_chars)?;
hits.push(hit);
if hits.len() >= k {
break;
}
}
tracing::debug!(rows = hits.len(), "kb-search lexical: search done");
Ok(hits)
}
fn index_version(&self) -> IndexVersion {
self.index_version.clone()
}
}
// ── Match-string construction ────────────────────────────────────────────
/// Translate a user-typed query into an FTS5 match string.
///
/// Rules (from the task spec):
///
/// - The query is wrapped in a single pair of `'...'` → strip the quotes
/// and pass the inner text through verbatim. The user has explicitly
/// opted into FTS5 syntax (e.g. `'rust AND cargo'`, `'foo*'`).
///
/// - Otherwise: split on whitespace, escape every token by wrapping it
/// in `"..."` (FTS5 string literal), with any inner `"` doubled. Join
/// with spaces — FTS5 default operator is implicit AND.
///
/// - An empty / whitespace-only token list → return `None` (caller
/// short-circuits to `Ok(vec![])`).
fn build_match_string(text: &str) -> Option<String> {
let trimmed = text.trim();
if trimmed.is_empty() {
return None;
}
if let Some(inner) = strip_single_quotes(trimmed) {
let inner_trim = inner.trim();
if inner_trim.is_empty() {
return None;
}
return Some(inner_trim.to_string());
}
let tokens: Vec<String> = trimmed
.split_whitespace()
.map(escape_fts5_token)
.collect();
if tokens.is_empty() {
None
} else {
Some(tokens.join(" "))
}
}
/// Return `Some(inner)` if `s` is wrapped in a matching pair of single
/// quotes (`'...'`), otherwise `None`. We require the closing quote to
/// be the last character so `'foo' bar` doesn't accidentally engage
/// raw-FTS5 mode.
fn strip_single_quotes(s: &str) -> Option<&str> {
let bytes = s.as_bytes();
if bytes.len() >= 2 && bytes[0] == b'\'' && bytes[bytes.len() - 1] == b'\'' {
Some(&s[1..s.len() - 1])
} else {
None
}
}
/// FTS5-escape one token by wrapping it in double quotes (FTS5 string
/// literal). Inner `"` are escaped by doubling per FTS5 grammar. This is
/// the simple-and-safe approach that defangs every special character —
/// `(`, `)`, `*`, `^`, `:`, `"`, etc. — without trying to parse FTS5
/// expressions.
fn escape_fts5_token(tok: &str) -> String {
let mut out = String::with_capacity(tok.len() + 2);
out.push('"');
for ch in tok.chars() {
if ch == '"' {
out.push('"');
out.push('"');
} else {
out.push(ch);
}
}
out.push('"');
out
}
// ── SQL execution ────────────────────────────────────────────────────────
/// Raw row shape mirroring the columns selected by [`run_query`]. Kept
/// internal — every public path constructs `SearchHit` from this.
struct RawRow {
chunk_id: String,
doc_id: String,
bm25_raw: f64,
snippet: String,
heading_path_json: String,
section_label: Option<String>,
source_spans_json: String,
chunker_version: String,
workspace_path: String,
}
/// Build + execute the FTS5 query. The SQL pattern is the one documented
/// in `tasks/p2/p2-2-lexical-retriever.md` (§Behavior contract).
fn run_query(
conn: &Connection,
match_str: &str,
snippet_words: usize,
filters: &SearchFilters,
fetch_limit: usize,
) -> Result<Vec<RawRow>> {
// Build the dynamic SQL + positional parameter vector. Positional `?`
// is used (not named bindings) because the dynamic IN-list for
// `tags_any` is most natural with `params_from_iter`.
let mut sql = String::from(
"SELECT \
f.chunk_id, f.doc_id, \
bm25(chunks_fts) AS score, \
snippet(chunks_fts, 3, '', '', '…', ?) AS snippet, \
c.heading_path_json, c.section_label, c.source_spans_json, \
c.chunker_version, \
d.workspace_path \
FROM chunks_fts f \
JOIN chunks c ON c.chunk_id = f.chunk_id \
JOIN documents d ON d.doc_id = f.doc_id",
);
let mut params: Vec<Box<dyn ToSql>> = Vec::new();
// 1) snippet word count.
params.push(Box::new(snippet_words as i64));
// 2) MATCH expression.
sql.push_str(" WHERE chunks_fts MATCH ?");
params.push(Box::new(match_str.to_owned()));
// tags_any: doc must own at least one of the requested tags.
if !filters.tags_any.is_empty() {
sql.push_str(
" AND f.doc_id IN (SELECT doc_id FROM document_tags WHERE tag IN (",
);
for (i, tag) in filters.tags_any.iter().enumerate() {
if i > 0 {
sql.push(',');
}
sql.push('?');
params.push(Box::new(tag.clone()));
}
sql.push_str("))");
}
if let Some(lang) = &filters.lang {
sql.push_str(" AND d.lang = ?");
params.push(Box::new(lang.0.clone()));
}
if let Some(trust_min) = &filters.trust_min {
// Mirror `kebab_store_sqlite::documents::list_documents` ranking:
// Generated < Secondary < Primary. Doing the rank in SQL
// (rather than post-filtering) keeps the row stream short
// when the workspace contains many low-trust docs.
sql.push_str(
" AND CASE d.trust_level \
WHEN 'primary' THEN 3 \
WHEN 'secondary' THEN 2 \
WHEN 'generated' THEN 1 \
ELSE 0 \
END >= ?",
);
let rank: i64 = match trust_min {
TrustLevel::Primary => 3,
TrustLevel::Secondary => 2,
TrustLevel::Generated => 1,
};
params.push(Box::new(rank));
}
// path_glob is intentionally NOT applied here — see module comment
// on PATH_GLOB_OVERFETCH and the post-filter in `LexicalRetriever::search`.
// Determinism: tie-break on chunk_id so equal bm25 scores produce a
// stable order across runs. `f.chunk_id` is the FTS row's UNINDEXED
// copy of the same value as `c.chunk_id`; either side works.
sql.push_str(" ORDER BY score, f.chunk_id LIMIT ?");
params.push(Box::new(i64::try_from(fetch_limit).unwrap_or(i64::MAX)));
let mut stmt = conn
.prepare(&sql)
.context("kb-search lexical: prepare FTS5 statement")?;
let rows = stmt
.query_map(params_from_iter(params.iter().map(|b| b.as_ref())), row_from_sql)
.context("kb-search lexical: execute FTS5 query")?;
let mut out: Vec<RawRow> = Vec::new();
for r in rows {
out.push(r.context("kb-search lexical: read row")?);
}
Ok(out)
}
fn row_from_sql(row: &Row<'_>) -> rusqlite::Result<RawRow> {
Ok(RawRow {
chunk_id: row.get(0)?,
doc_id: row.get(1)?,
bm25_raw: row.get(2)?,
snippet: row.get(3)?,
heading_path_json: row.get(4)?,
section_label: row.get(5)?,
source_spans_json: row.get(6)?,
chunker_version: row.get(7)?,
workspace_path: row.get(8)?,
})
}
// ── Hit construction ─────────────────────────────────────────────────────
fn build_hit(
raw: RawRow,
rank: u32,
index_version: &IndexVersion,
snippet_chars: usize,
) -> Result<SearchHit> {
let normalized = normalize_bm25(raw.bm25_raw);
let heading_path: Vec<String> = serde_json::from_str(&raw.heading_path_json)
.context("kb-search lexical: deserialize heading_path_json")?;
let source_spans: Vec<SourceSpan> = serde_json::from_str(&raw.source_spans_json)
.context("kb-search lexical: deserialize source_spans_json")?;
let workspace_path = WorkspacePath::new(raw.workspace_path)
.context("kb-search lexical: documents.workspace_path violates WorkspacePath invariant")?;
let citation = citation_from_first_span(
&raw.chunk_id,
workspace_path.clone(),
raw.section_label.clone(),
source_spans.first(),
);
// FTS5's snippet() respects the word budget but produces a
// character-length we can't predict precisely (token boundaries vary
// with the tokenizer). The contract caps at `snippet_chars`; trim
// defensively if SQLite ever returns a longer string.
let snippet = trim_snippet(&raw.snippet, snippet_chars);
Ok(SearchHit {
rank,
chunk_id: ChunkId(raw.chunk_id),
doc_id: DocumentId(raw.doc_id),
doc_path: workspace_path,
heading_path,
section_label: raw.section_label,
snippet,
citation,
retrieval: RetrievalDetail {
method: SearchMode::Lexical,
fusion_score: normalized,
lexical_score: Some(normalized),
vector_score: None,
lexical_rank: Some(rank),
vector_rank: None,
},
index_version: index_version.clone(),
embedding_model: None,
chunker_version: ChunkerVersion(raw.chunker_version),
})
}
/// Map the raw bm25 score (FTS5 returns a *negative* number; lower is
/// better) into a positive score in `(0, 1]`. The formula
/// `score = -bm25 / (1 + |bm25|)` is monotonic, smooth, and bounded —
/// suitable both for human display and for use as an RRF input.
fn normalize_bm25(bm25_raw: f64) -> f32 {
let abs = bm25_raw.abs();
let normalized = -bm25_raw / (1.0_f64 + abs);
normalized as f32
}
/// Cap the snippet at `max_chars` characters (Unicode scalar values, not
/// bytes — matches the §6.4 setting's "characters" semantics). Returns
/// the input unchanged when already short enough.
fn trim_snippet(s: &str, max_chars: usize) -> String {
// We slice on Unicode scalar values per §6.4's "characters" semantics; this
// can orphan a combining mark in extreme cases (Hebrew niqqud, Devanagari)
// but matches the spec's char-budget definition.
if s.chars().count() <= max_chars {
return s.to_string();
}
s.chars().take(max_chars).collect()
}
// ── path_glob ────────────────────────────────────────────────────────────
/// Compile a `path_glob` pattern. We enable `literal_separator` so `*`
/// does NOT cross `/` — design requires `*` to match within a single
/// path segment, not across them. (`globset`'s default is to let `*`
/// span separators.)
fn compile_glob(pattern: &str) -> Result<GlobMatcher> {
let g = globset::GlobBuilder::new(pattern)
.literal_separator(true)
.build()
.with_context(|| format!("kb-search lexical: invalid path_glob {pattern:?}"))?;
Ok(g.compile_matcher())
}
// ── Unit tests for pure helpers ──────────────────────────────────────────
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn build_match_string_empty_returns_none() {
assert!(build_match_string("").is_none());
assert!(build_match_string(" ").is_none());
assert!(build_match_string("''").is_none());
assert!(build_match_string("' '").is_none());
}
#[test]
fn build_match_string_default_is_quoted_and_anded() {
let s = build_match_string("rust cargo").unwrap();
// Two tokens, each quoted, joined by a space (implicit AND).
assert_eq!(s, r#""rust" "cargo""#);
}
#[test]
fn build_match_string_escapes_special_chars() {
// `*`, `(`, `)`, `:`, `^`, `"` should all be wrapped inside
// FTS5 string-literal quotes so they're treated as literal
// text rather than FTS5 operators.
let s = build_match_string(r#"foo* (bar) baz:qux ^head he"llo"#).unwrap();
assert_eq!(
s,
r#""foo*" "(bar)" "baz:qux" "^head" "he""llo""#
);
// The doubled `""` is FTS5's way of embedding a literal quote
// inside a string literal.
assert!(s.contains(r#"he""llo"#));
// Sanity: every special character lives between matching `"`
// delimiters — there is no bare-token (unquoted) span anywhere.
// We check this by confirming the string starts and ends with `"`
// and the count of unescaped `"` is even (each token is wrapped).
assert!(s.starts_with('"') && s.ends_with('"'));
}
#[test]
fn build_match_string_passthrough_when_single_quoted() {
// The FTS5 expression is preserved verbatim.
let s = build_match_string("'foo OR bar*'").unwrap();
assert_eq!(s, "foo OR bar*");
}
#[test]
fn normalize_bm25_top_score_in_unit_interval() {
// A "perfect" hit is bm25 = -1.0 → normalized 0.5.
// A high-relevance hit (bm25 = -10.0) → 10/11 ≈ 0.909.
let high = normalize_bm25(-10.0);
assert!(high > 0.0 && high <= 1.0, "got {high}");
let medium = normalize_bm25(-1.0);
assert!((medium - 0.5).abs() < 1e-6);
}
#[test]
fn normalize_bm25_monotonic() {
// Lower (more-negative) bm25 must map to a higher normalized score.
let a = normalize_bm25(-2.0);
let b = normalize_bm25(-1.0);
assert!(a > b, "{a} should exceed {b}");
}
#[test]
fn trim_snippet_caps_at_char_count() {
let s = "a".repeat(300);
let trimmed = trim_snippet(&s, 220);
assert_eq!(trimmed.chars().count(), 220);
}
#[test]
fn trim_snippet_passthrough_when_short() {
let s = "short";
assert_eq!(trim_snippet(s, 220), "short");
}
#[test]
fn build_citation_line_round_trip() {
use kebab_core::Citation;
let p = WorkspacePath::new("a/b.md".to_string()).unwrap();
let span = SourceSpan::Line { start: 7, end: 12 };
let c = citation_from_first_span("c1", p.clone(), Some("S1".to_string()), Some(&span));
match c {
Citation::Line {
start,
end,
ref section,
path: ref pp,
} => {
assert_eq!(start, 7);
assert_eq!(end, 12);
assert_eq!(section.as_deref(), Some("S1"));
assert_eq!(pp, &p);
}
other => panic!("expected Citation::Line, got {other:?}"),
}
}
#[test]
fn build_citation_page_forwards_section() {
use kebab_core::Citation;
let p = WorkspacePath::new("doc.pdf".to_string()).unwrap();
let span = SourceSpan::Page {
page: 4,
char_start: None,
char_end: None,
};
let c = citation_from_first_span("c1", p, Some("Intro".to_string()), Some(&span));
match c {
Citation::Page {
page,
ref section,
..
} => {
assert_eq!(page, 4);
assert_eq!(section.as_deref(), Some("Intro"));
}
other => panic!("expected Citation::Page, got {other:?}"),
}
}
#[test]
fn build_citation_none_falls_back_to_line_one() {
use kebab_core::Citation;
let p = WorkspacePath::new("x.md".to_string()).unwrap();
let c = citation_from_first_span("c1", p, None, None);
match c {
Citation::Line { start, end, .. } => {
assert_eq!((start, end), (1, 1));
}
other => panic!("expected fallback Citation::Line, got {other:?}"),
}
}
#[test]
fn compile_glob_rejects_invalid_pattern() {
// `[` is a character-class opener; an unclosed class is invalid.
let r = compile_glob("notes/[abc");
assert!(r.is_err());
}
#[test]
fn compile_glob_star_does_not_cross_slash() {
// This is the design invariant: `*` must NOT match `/`.
let m = compile_glob("notes/*.md").unwrap();
assert!(m.is_match("notes/foo.md"));
assert!(!m.is_match("notes/sub/foo.md"));
}
}

26
crates/kebab-search/src/lib.rs Normal file
View File

@@ -0,0 +1,26 @@
//! `kb-search` — `kebab_core::Retriever` implementations.
//!
//! - [`LexicalRetriever`] (P2-2): SQLite-FTS5 + bm25 backed retriever
//! for `SearchMode::Lexical`.
//! - [`VectorRetriever`] (P3-4): wraps a `dyn VectorStore` (typically
//! `kb-store-vector::LanceVectorStore`) and a `dyn Embedder`,
//! hydrating SQLite metadata for full `SearchHit`s.
//! - [`HybridRetriever`] (P3-4): composes lexical + vector retrievers,
//! dispatches by `SearchMode`, fuses Hybrid via [`FusionPolicy::Rrf`].
//!
//! Allowed deps per the P2-2 + P3-4 task specs: `kb-core`, `kb-config`,
//! `kb-store-sqlite`, `kb-store-vector`, `kb-embed` (trait re-export
//! only — concrete adapters like `kb-embed-local` are runtime-injected
//! via `Arc<dyn Embedder>`), `rusqlite`, `globset`, `serde_json`,
//! `tracing`, `thiserror`, `anyhow`. Forbidden: `kb-source-fs`,
//! `kb-parse-md`, `kb-normalize`, `kb-chunk`, `kb-embed-local` (concrete
//! adapter), `kb-llm*`, `kb-rag`, `kb-tui`, `kb-desktop`.
mod citation_helper;
mod hybrid;
mod lexical;
mod vector;
pub use hybrid::{FusionPolicy, HybridRetriever};
pub use lexical::LexicalRetriever;
pub use vector::VectorRetriever;

338
crates/kebab-search/src/vector.rs Normal file
View File

@@ -0,0 +1,338 @@
//! Vector retriever — design §3.7 / §7.2 / §1.6.
//!
//! Wraps a `dyn VectorStore` + `dyn Embedder` + the SQLite metadata
//! store into a `kebab_core::Retriever`. The vector store knows how to
//! find the nearest chunks by cosine on the embedding column; SQLite
//! owns the human-readable metadata (heading_path / section_label /
//! source_spans / chunker_version / workspace_path) needed for
//! `SearchHit` and `Citation`. The retriever stitches them together
//! per spec §7.2.
//!
//! Snippet policy: this retriever has no FTS5 highlighter to lean on,
//! so the `snippet` field is the chunk text trimmed to
//! `config.search.snippet_chars` Unicode scalar values. The lexical
//! retriever does query-token highlighting; downstream UI code should
//! continue to surface lexical snippets for hybrid hits where the
//! lexical side contributed (handled in `HybridRetriever::search`).
use std::collections::HashMap;
use std::sync::Arc;
use anyhow::{Context, Result};
use kebab_core::{
ChunkId, ChunkerVersion, DocumentId, Embedder, EmbeddingInput, EmbeddingKind,
IndexVersion, RetrievalDetail, Retriever, SearchHit, SearchMode, SearchQuery,
SourceSpan, VectorHit, VectorStore, WorkspacePath,
};
use kebab_store_sqlite::SqliteStore;
use rusqlite::params_from_iter;
use crate::citation_helper::citation_from_first_span;
/// Default `k` when `SearchQuery::k == 0`. Mirrors §6.4 default_k=10
/// and the lexical retriever's `DEFAULT_K`.
const DEFAULT_K: usize = 10;
/// Over-fetch multiplier passed to `VectorStore::search` so that
/// SQLite-side filter losses (tags / lang / trust / path_glob) still
/// leave at least `k` candidates. The Lance store already applies the
/// same filters internally; the extra `* 2` is the spec-mandated
/// safety margin for the `Retriever` layer (§7.2 spec line 138).
const VECTOR_OVERFETCH_MULTIPLIER: usize = 2;
/// Wraps a vector store + embedder into a [`Retriever`].
///
/// `VectorStore` is not declared `Send + Sync` in `kb-core::traits`,
/// but `Retriever` requires both. We constrain the trait objects
/// here so callers must hand us implementations that already are
/// (`LanceVectorStore` is `Send + Sync` thanks to its
/// `Connection`/`Runtime` ownership; the trait is sync-method-only).
pub struct VectorRetriever {
store: Arc<dyn VectorStore + Send + Sync>,
embed: Arc<dyn Embedder>,
sqlite: Arc<SqliteStore>,
index_version: IndexVersion,
snippet_chars: usize,
}
impl VectorRetriever {
/// Construct with `index_version` derived from the configured
/// embedding model + dimensions, and snippet width pulled from
/// `kb-config`'s defaults.
///
/// The explicit `index_version` form is [`Self::with_settings`].
pub fn new(
store: Arc<dyn VectorStore + Send + Sync>,
embed: Arc<dyn Embedder>,
sqlite: Arc<SqliteStore>,
index_version: IndexVersion,
) -> Self {
let cfg = kebab_config::Config::defaults();
Self::with_settings(store, embed, sqlite, index_version, cfg.search.snippet_chars)
}
/// Construct with explicit `snippet_chars`. Mirrors the lexical
/// retriever's `with_settings` constructor for callers that have
/// already loaded a `Config`.
pub fn with_settings(
store: Arc<dyn VectorStore + Send + Sync>,
embed: Arc<dyn Embedder>,
sqlite: Arc<SqliteStore>,
index_version: IndexVersion,
snippet_chars: usize,
) -> Self {
Self {
store,
embed,
sqlite,
index_version,
snippet_chars,
}
}
}
impl Retriever for VectorRetriever {
fn search(&self, query: &SearchQuery) -> Result<Vec<SearchHit>> {
let k = if query.k == 0 { DEFAULT_K } else { query.k };
tracing::debug!(
text_len = query.text.len(),
k,
"kb-search vector: search start"
);
// Empty / whitespace-only queries — short-circuit. The
// embedder would still produce a vector for an empty string,
// but nearest-neighbours on the centroid of "" is meaningless
// and only forces a wasted Lance scan.
if query.text.trim().is_empty() {
return Ok(Vec::new());
}
// 1. Embed the query as `Query` kind (e5-style asymmetry —
// documents and queries have different prefixes).
let inputs = [EmbeddingInput {
text: &query.text,
kind: EmbeddingKind::Query,
}];
let mut embeddings = self
.embed
.embed(&inputs)
.context("kb-search vector: embed query")?;
if embeddings.len() != 1 {
anyhow::bail!(
"kb-search vector: embedder returned {} vectors for one input",
embeddings.len()
);
}
let query_vec = embeddings.remove(0);
// 2. Over-fetch from the vector store. The Lance store
// applies `filter_chunks` internally, so we pass `query.filters`
// through and trust the post-filter pass to honour them.
// `saturating_mul(2)` is always ≥ k for any usize k, so we
// don't need an extra `.max(k)` clamp.
let overfetch = k.saturating_mul(VECTOR_OVERFETCH_MULTIPLIER);
let raw_hits = self
.store
.search(&query_vec, overfetch, &query.filters)
.context("kb-search vector: VectorStore::search")?;
if raw_hits.is_empty() {
tracing::debug!("kb-search vector: store returned no hits");
return Ok(Vec::new());
}
// 3. Hydrate metadata from SQLite for the candidate ids in
// one round-trip. Order is preserved by the caller via the
// HashMap lookup at hit-construction time.
let candidate_ids: Vec<&str> =
raw_hits.iter().map(|h| h.chunk_id.0.as_str()).collect();
let hydration = hydrate_chunks(&self.sqlite, &candidate_ids)
.context("kb-search vector: hydrate chunk metadata")?;
// 4. Build `SearchHit` for the first `k` raw hits that pass
// hydration (a missing row would be a filter-induced drop —
// Lance returned the chunk but SQLite filtered it out, or
// the chunk was deleted between Lance's read and ours).
let model_id = self.embed.model_id();
let mut hits: Vec<SearchHit> = Vec::with_capacity(k.min(raw_hits.len()));
let mut rank: u32 = 0;
for hit in raw_hits {
let Some(meta) = hydration.get(hit.chunk_id.0.as_str()) else {
continue;
};
rank = rank.saturating_add(1);
hits.push(build_hit(
hit,
meta,
rank,
&self.index_version,
&model_id,
self.snippet_chars,
)?);
if hits.len() >= k {
break;
}
}
tracing::debug!(rows = hits.len(), "kb-search vector: search done");
Ok(hits)
}
fn index_version(&self) -> IndexVersion {
self.index_version.clone()
}
}
// ── Hydration ────────────────────────────────────────────────────────────
/// Subset of `chunks` + `documents` metadata needed to build a
/// `SearchHit` from a `VectorHit`. Pulled in one round-trip so the
/// per-hit construction loop stays O(1) per row.
struct ChunkMeta {
text: String,
heading_path_json: String,
section_label: Option<String>,
source_spans_json: String,
chunker_version: String,
doc_id: String,
workspace_path: String,
}
fn hydrate_chunks(
sqlite: &SqliteStore,
chunk_ids: &[&str],
) -> Result<HashMap<String, ChunkMeta>> {
if chunk_ids.is_empty() {
return Ok(HashMap::new());
}
// Deduplicate the IN-list — Lance can repeat a chunk_id across
// batches in pathological cases. A HashMap key dedupes in the
// result anyway, but keeping the placeholder count tight is good
// hygiene.
let mut seen = std::collections::HashSet::new();
let unique: Vec<&str> = chunk_ids
.iter()
.copied()
.filter(|id| seen.insert(*id))
.collect();
let placeholders = vec!["?"; unique.len()].join(",");
let sql = format!(
"SELECT \
c.chunk_id, c.text, c.heading_path_json, c.section_label, \
c.source_spans_json, c.chunker_version, \
c.doc_id, d.workspace_path \
FROM chunks c \
JOIN documents d ON d.doc_id = c.doc_id \
WHERE c.chunk_id IN ({placeholders})"
);
let conn = sqlite.read_conn();
let mut stmt = conn
.prepare(&sql)
.context("kb-search vector: prepare hydration statement")?;
let rows = stmt
.query_map(
// `unique` is a `Vec<&str>`; `&str` implements `ToSql`
// directly, so we hand the iterator straight to
// `params_from_iter` without copying.
params_from_iter(unique.iter().copied()),
|row| {
let chunk_id: String = row.get(0)?;
Ok((
chunk_id,
ChunkMeta {
text: row.get(1)?,
heading_path_json: row.get(2)?,
section_label: row.get(3)?,
source_spans_json: row.get(4)?,
chunker_version: row.get(5)?,
doc_id: row.get(6)?,
workspace_path: row.get(7)?,
},
))
},
)
.context("kb-search vector: execute hydration query")?;
let mut out: HashMap<String, ChunkMeta> = HashMap::with_capacity(unique.len());
for row in rows {
let (chunk_id, meta) =
row.context("kb-search vector: read hydration row")?;
out.insert(chunk_id, meta);
}
Ok(out)
}
fn build_hit(
hit: VectorHit,
meta: &ChunkMeta,
rank: u32,
index_version: &IndexVersion,
model_id: &kebab_core::EmbeddingModelId,
snippet_chars: usize,
) -> Result<SearchHit> {
let heading_path: Vec<String> = serde_json::from_str(&meta.heading_path_json)
.context("kb-search vector: deserialize heading_path_json")?;
let source_spans: Vec<SourceSpan> = serde_json::from_str(&meta.source_spans_json)
.context("kb-search vector: deserialize source_spans_json")?;
let workspace_path = WorkspacePath::new(meta.workspace_path.clone()).context(
"kb-search vector: documents.workspace_path violates WorkspacePath invariant",
)?;
let citation = citation_from_first_span(
&hit.chunk_id.0,
workspace_path.clone(),
meta.section_label.clone(),
source_spans.first(),
);
let snippet = trim_snippet(&meta.text, snippet_chars);
let score = hit.score;
Ok(SearchHit {
rank,
chunk_id: ChunkId(hit.chunk_id.0),
doc_id: DocumentId(meta.doc_id.clone()),
doc_path: workspace_path,
heading_path,
section_label: meta.section_label.clone(),
snippet,
citation,
retrieval: RetrievalDetail {
method: SearchMode::Vector,
fusion_score: score,
lexical_score: None,
vector_score: Some(score),
lexical_rank: None,
vector_rank: Some(rank),
},
index_version: index_version.clone(),
embedding_model: Some(model_id.clone()),
chunker_version: ChunkerVersion(meta.chunker_version.clone()),
})
}
/// Cap the snippet at `max_chars` Unicode scalar values. Mirrors
/// `lexical::trim_snippet` so the two retrievers produce identically
/// shaped snippets for hybrid output.
fn trim_snippet(s: &str, max_chars: usize) -> String {
if s.chars().count() <= max_chars {
return s.to_string();
}
s.chars().take(max_chars).collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn trim_snippet_caps_at_char_count() {
let s = "a".repeat(300);
assert_eq!(trim_snippet(&s, 220).chars().count(), 220);
}
#[test]
fn trim_snippet_passthrough_when_short() {
assert_eq!(trim_snippet("short", 220), "short");
}
}

216
crates/kebab-search/tests/common/mod.rs Normal file
View File

@@ -0,0 +1,216 @@
//! Shared scaffolding for kb-search hybrid integration tests.
//!
//! # Test policy
//!
//! Integration tests in `hybrid.rs` that touch `LanceVectorStore`
//! are marked `#[ignore]` AND call [`require_avx_or_panic`] inside
//! the test body so a `--ignored` invocation on a non-AVX host
//! fails loudly with a clear message rather than crashing later
//! inside Lance's f32 SIMD kernel with `SIGILL`.
//!
//! See `crates/kb-store-vector/tests/common/mod.rs` for the
//! original P3-3 rationale; this is a copy because that crate's
//! test commons are test-only and not part of its public surface.
#![allow(dead_code)]
use std::sync::Arc;
use kebab_config::Config;
use kebab_core::{
ChunkId, DocumentId, EmbeddingId, EmbeddingInput, EmbeddingKind,
EmbeddingModelId, EmbeddingVersion, IndexVersion, VectorRecord, VectorStore,
};
use kebab_embed::{Embedder, MockEmbedder};
use kebab_search::{LexicalRetriever, VectorRetriever};
use kebab_store_sqlite::SqliteStore;
use kebab_store_vector::LanceVectorStore;
use rusqlite::params;
use tempfile::TempDir;
/// Panic if the host CPU lacks AVX. Called from every `#[ignore]`-d
/// integration test body so that `cargo test -- --ignored` on a
/// non-AVX host fails loudly with a clear message instead of crashing
/// later inside a Lance SIMD kernel with `SIGILL`.
pub fn require_avx_or_panic() {
#[cfg(target_arch = "x86_64")]
{
if !std::is_x86_feature_detected!("avx") {
panic!(
"kb-search hybrid integration test requires AVX-capable hardware; \
host CPU lacks AVX. Run on an AVX-capable machine."
);
}
}
}
/// Index version label used by hybrid integration tests so the
/// `index_version()` composite token is predictable in snapshots.
pub const TEST_LEX_INDEX_VERSION: &str = "v1.0-lex";
pub const TEST_VEC_INDEX_VERSION: &str = "v1.0-vec";
/// Embedding dimensions for tests. Kept small so MockEmbedder runs
/// fast and the Lance table stays compact on disk; production uses
/// 384 (multilingual-e5-small) but the retriever code is dim-agnostic.
pub const TEST_DIMENSIONS: usize = 16;
pub const TEST_MODEL_ID: &str = "mock-e5";
pub struct HybridEnv {
pub temp: TempDir,
pub config: Config,
pub sqlite: Arc<SqliteStore>,
pub vector_store: Arc<LanceVectorStore>,
pub embedder: Arc<MockEmbedder>,
}
impl HybridEnv {
pub fn new() -> Self {
let temp = tempfile::tempdir().expect("tempdir");
let mut config = Config::defaults();
config.storage.data_dir = temp.path().to_string_lossy().into_owned();
let sqlite = SqliteStore::open(&config).unwrap();
sqlite.run_migrations().unwrap();
let sqlite = Arc::new(sqlite);
let vector_store =
Arc::new(LanceVectorStore::new(&config, sqlite.clone()).unwrap());
let embedder = Arc::new(MockEmbedder::new(
EmbeddingModelId(TEST_MODEL_ID.to_string()),
EmbeddingVersion("v1".to_string()),
TEST_DIMENSIONS,
));
Self {
temp,
config,
sqlite,
vector_store,
embedder,
}
}
/// Build a `LexicalRetriever` over the shared SQLite store.
pub fn lexical_retriever(&self) -> LexicalRetriever {
LexicalRetriever::new(
Arc::clone(&self.sqlite),
IndexVersion(TEST_LEX_INDEX_VERSION.to_string()),
)
}
/// Build a `VectorRetriever` over the shared LanceVectorStore +
/// MockEmbedder + SQLite store.
pub fn vector_retriever(&self) -> VectorRetriever {
let store: Arc<dyn VectorStore + Send + Sync> =
Arc::clone(&self.vector_store) as Arc<dyn VectorStore + Send + Sync>;
let embed: Arc<dyn Embedder> =
Arc::clone(&self.embedder) as Arc<dyn Embedder>;
VectorRetriever::new(
store,
embed,
Arc::clone(&self.sqlite),
IndexVersion(TEST_VEC_INDEX_VERSION.to_string()),
)
}
/// Insert (asset, document, document_tags, chunk) rows directly.
/// We seed without going through `DocumentStore::put_document`
/// to keep this crate's test deps inside the Allowed list (no
/// `kb-parse-md` / `kb-normalize` / `kb-chunk`). The `chunks` row
/// also fires the V002 FTS5 triggers, so the lexical retriever
/// can find the row by `MATCH` without a manual rebuild.
pub fn seed_chunk(
&self,
chunk_id: &str,
doc_id: &str,
workspace_path: &str,
text: &str,
heading_path: &[&str],
tags: &[&str],
) {
let asset_id = format!("a{}", &doc_id[..31]);
let conn = self.sqlite.read_conn();
conn.execute(
"INSERT OR IGNORE INTO assets (
asset_id, source_uri, workspace_path, media_type, byte_len,
checksum, storage_kind, storage_path, discovered_at
) VALUES (?, ?, ?, '\"markdown\"', 0,
'deadbeefdeadbeefdeadbeefdeadbeef',
'reference', ?, '1970-01-01T00:00:00Z')",
params![
asset_id,
format!("file://{workspace_path}"),
workspace_path,
workspace_path,
],
)
.unwrap();
conn.execute(
"INSERT OR IGNORE INTO documents (
doc_id, asset_id, workspace_path, title, lang, source_type,
trust_level, parser_version, doc_version, schema_version,
metadata_json, provenance_json, created_at, updated_at
) VALUES (?, ?, ?, NULL, 'en', 'markdown', 'primary', 'v1', 1, 1,
'{}', '{}', '1970-01-01T00:00:00Z', '1970-01-01T00:00:00Z')",
params![doc_id, asset_id, workspace_path],
)
.unwrap();
for t in tags {
conn.execute(
"INSERT OR IGNORE INTO document_tags (doc_id, tag) VALUES (?, ?)",
params![doc_id, t],
)
.unwrap();
}
let heading_json = serde_json::to_string(heading_path).unwrap();
conn.execute(
"INSERT OR IGNORE INTO chunks (
chunk_id, doc_id, text, heading_path_json, section_label,
source_spans_json, token_estimate, chunker_version,
policy_hash, block_ids_json, created_at
) VALUES (?, ?, ?, ?, NULL,
'[{\"kind\":\"line\",\"start\":1,\"end\":3}]',
1, 'v1', 'h', '[]', '1970-01-01T00:00:00Z')",
params![chunk_id, doc_id, text, heading_json],
)
.unwrap();
}
/// Embed `text` as a Document and upsert it as the embedding for
/// `chunk_id`. Drives the same code path production uses:
/// MockEmbedder → VectorRecord → LanceVectorStore::upsert →
/// embedding_records committed.
pub fn embed_and_upsert(
&self,
chunk_id: &str,
doc_id: &str,
text: &str,
heading_path: &[&str],
) {
let inputs = [EmbeddingInput {
text,
kind: EmbeddingKind::Document,
}];
let mut vecs = self.embedder.embed(&inputs).unwrap();
let vector = vecs.remove(0);
let record = VectorRecord {
chunk_id: ChunkId(chunk_id.to_string()),
embedding_id: EmbeddingId(format!("e{}", &chunk_id[..31])),
vector,
doc_id: DocumentId(doc_id.to_string()),
text: text.to_string(),
heading_path: heading_path.iter().map(|s| s.to_string()).collect(),
model_id: EmbeddingModelId(TEST_MODEL_ID.to_string()),
model_version: EmbeddingVersion("v1".to_string()),
dimensions: TEST_DIMENSIONS,
};
self.vector_store.upsert(&[record]).unwrap();
}
}
/// Pad a short prefix to the 32-hex shape `kebab_core` newtypes expect.
pub fn id32(prefix: &str) -> String {
let mut s = prefix.to_string();
while s.len() < 32 {
s.push('0');
}
s.truncate(32);
s
}

42
crates/kebab-search/tests/fixtures/search/hybrid/run-1.json vendored Normal file
View File

@@ -0,0 +1,42 @@
[
{
"chunk_id": "c1000000000000000000000000000000",
"fusion_score_positive": true,
"lex_some": true,
"lexical_rank": 1,
"method": "hybrid",
"rank": 1,
"vec_some": true,
"vector_rank": 3
},
{
"chunk_id": "c2000000000000000000000000000000",
"fusion_score_positive": true,
"lex_some": true,
"lexical_rank": 2,
"method": "hybrid",
"rank": 2,
"vec_some": true,
"vector_rank": 2
},
{
"chunk_id": "c4000000000000000000000000000000",
"fusion_score_positive": true,
"lex_some": false,
"lexical_rank": null,
"method": "hybrid",
"rank": 3,
"vec_some": true,
"vector_rank": 1
},
{
"chunk_id": "c3000000000000000000000000000000",
"fusion_score_positive": true,
"lex_some": false,
"lexical_rank": null,
"method": "hybrid",
"rank": 4,
"vec_some": true,
"vector_rank": 4
}
]

60
crates/kebab-search/tests/fixtures/search/lexical/run-1.json vendored Normal file
View File

@@ -0,0 +1,60 @@
[
{
"chunk_id": "c3000000000000000000000000000000",
"chunker_version": "v1",
"citation": {
"end": 8,
"kind": "line",
"path": "notes/snap.md",
"section": "Snap",
"start": 7
},
"doc_id": "d0000000000000000000000000000000",
"doc_path": "notes/snap.md",
"embedding_model": null,
"heading_path": [
"Snap"
],
"index_version": "v1.0",
"rank": 1,
"retrieval": {
"fusion_score": 1.4490997273242101e-6,
"lexical_rank": 1,
"lexical_score": 1.4490997273242101e-6,
"method": "lexical",
"vector_rank": null,
"vector_score": null
},
"section_label": "Snap",
"snippet": "alpha alpha"
},
{
"chunk_id": "c1000000000000000000000000000000",
"chunker_version": "v1",
"citation": {
"end": 2,
"kind": "line",
"path": "notes/snap.md",
"section": "Snap",
"start": 1
},
"doc_id": "d0000000000000000000000000000000",
"doc_path": "notes/snap.md",
"embedding_model": null,
"heading_path": [
"Snap"
],
"index_version": "v1.0",
"rank": 2,
"retrieval": {
"fusion_score": 9.641424867368187e-7,
"lexical_rank": 2,
"lexical_score": 9.641424867368187e-7,
"method": "lexical",
"vector_rank": null,
"vector_score": null
},
"section_label": "Snap",
"snippet": "alpha bravo charlie"
}
]

213
crates/kebab-search/tests/hybrid.rs Normal file
View File

@@ -0,0 +1,213 @@
//! Hybrid integration tests — touch a real `LanceVectorStore` +
//! `SqliteStore` + `MockEmbedder`. These tests are `#[ignore]`-d and
//! AVX-gated; see `tests/common/mod.rs` for the policy rationale.
//!
//! Mock-retriever unit tests live alongside the implementation in
//! `crates/kb-search/src/hybrid.rs` (no Lance, no AVX needed) — the
//! tests here exercise the full plumbing with the real Lance store.
mod common;
use std::path::PathBuf;
use std::sync::Arc;
use common::{
HybridEnv, id32, require_avx_or_panic, TEST_LEX_INDEX_VERSION, TEST_VEC_INDEX_VERSION,
};
use kebab_core::{
Retriever, SearchFilters, SearchHit, SearchMode, SearchQuery,
};
use kebab_search::{FusionPolicy, HybridRetriever};
use serde_json::json;
fn build_hybrid(env: &HybridEnv) -> HybridRetriever {
let lex: Arc<dyn Retriever> = Arc::new(env.lexical_retriever());
let vec: Arc<dyn Retriever> = Arc::new(env.vector_retriever());
HybridRetriever::with_policy(lex, vec, FusionPolicy::Rrf { k_rrf: 60 }, 5)
}
/// Seed a tiny corpus that lets us prove hybrid recall ≥ each side
/// independently. Two chunks are lexical-only matches ("rust cargo");
/// two chunks are vector-only matches (their text doesn't contain
/// the query token but their embedding still scores nearby because
/// MockEmbedder's hash distributes over all chunks).
fn seed_disjoint_corpus(env: &HybridEnv) -> Vec<String> {
// The lexical side will only match chunks that contain the query
// tokens. The vector side will rank ALL chunks by embedding
// similarity to the query — even ones whose text doesn't share
// a token with the query.
let chunks = [
// (chunk_id, doc_id, path, text, headings)
(id32("c1"), id32("d1"), "notes/rust1.md", "rust cargo macros", &["A"][..]),
(id32("c2"), id32("d2"), "notes/rust2.md", "rust traits and lifetimes", &["B"][..]),
(id32("c3"), id32("d3"), "notes/python.md", "python dataclasses tutorial", &["C"][..]),
(id32("c4"), id32("d4"), "notes/go.md", "go interfaces and channels", &["D"][..]),
];
let mut ids = Vec::new();
for (cid, did, path, text, headings) in &chunks {
env.seed_chunk(cid, did, path, text, headings, &[]);
env.embed_and_upsert(cid, did, text, headings);
ids.push(cid.clone());
}
ids
}
#[test]
#[ignore = "requires AVX-capable hardware (LanceDB)"]
fn hybrid_recall_disjoint_returns_union() {
require_avx_or_panic();
let env = HybridEnv::new();
let _ids = seed_disjoint_corpus(&env);
let h = build_hybrid(&env);
let q = SearchQuery {
text: "rust".to_string(),
mode: SearchMode::Hybrid,
k: 4,
filters: SearchFilters::default(),
};
let hits = h.search(&q).unwrap();
// The vector side will return up to 4 candidates regardless of
// text overlap; the lexical side will return only the rust* ones.
// Together the union must cover at least the lexical hits AND
// include at least one non-lexical chunk if vector found one.
assert!(!hits.is_empty(), "hybrid must return at least one hit");
// Every hit's RetrievalDetail.method must be Hybrid.
for h in &hits {
assert_eq!(h.retrieval.method, SearchMode::Hybrid);
// At least one of lex/vec_score must be Some.
assert!(
h.retrieval.lexical_score.is_some() || h.retrieval.vector_score.is_some(),
"hybrid hit must carry at least one mode's score"
);
}
// index_version composite token.
let iv = h.index_version();
assert!(iv.0.starts_with("hybrid:"));
assert!(iv.0.contains(TEST_LEX_INDEX_VERSION));
assert!(iv.0.contains(TEST_VEC_INDEX_VERSION));
// Lexical-only chunks (c1, c2) MUST appear: they're the only ones
// matching the FTS5 query, and the vector side over-fetches enough
// to include them too.
let ids: Vec<&str> = hits.iter().map(|h| h.chunk_id.0.as_str()).collect();
assert!(ids.contains(&id32("c1").as_str()));
assert!(ids.contains(&id32("c2").as_str()));
}
#[test]
#[ignore = "requires AVX-capable hardware (LanceDB)"]
fn hybrid_determinism_same_query_twice() {
require_avx_or_panic();
let env = HybridEnv::new();
let _ = seed_disjoint_corpus(&env);
let h = build_hybrid(&env);
let q = SearchQuery {
text: "rust".to_string(),
mode: SearchMode::Hybrid,
k: 4,
filters: SearchFilters::default(),
};
let a = h.search(&q).unwrap();
let b = h.search(&q).unwrap();
assert_eq!(a, b, "identical query must yield byte-identical Vec<SearchHit>");
}
#[test]
#[ignore = "requires AVX-capable hardware (LanceDB)"]
fn hybrid_snapshot_run_1() {
require_avx_or_panic();
let env = HybridEnv::new();
let _ = seed_disjoint_corpus(&env);
let h = build_hybrid(&env);
let q = SearchQuery {
text: "rust".to_string(),
mode: SearchMode::Hybrid,
k: 4,
filters: SearchFilters::default(),
};
let hits = h.search(&q).unwrap();
// Snapshot pins the structural shape:
// - chunk_id ordering
// - which side contributed (lexical_rank / vector_rank
// populated as Some/None)
// - that fusion_score is non-increasing
// - method = Hybrid for every hit
let actual = json!(
hits.iter().map(|h: &SearchHit| json!({
"chunk_id": h.chunk_id.0,
"rank": h.rank,
"method": h.retrieval.method,
"lexical_rank": h.retrieval.lexical_rank,
"vector_rank": h.retrieval.vector_rank,
"lex_some": h.retrieval.lexical_score.is_some(),
"vec_some": h.retrieval.vector_score.is_some(),
"fusion_score_positive": h.retrieval.fusion_score > 0.0,
})).collect::<Vec<_>>()
);
let fixture = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("tests")
.join("fixtures")
.join("search")
.join("hybrid")
.join("run-1.json");
if std::env::var_os("KB_UPDATE_SNAPSHOTS").is_some() {
std::fs::create_dir_all(fixture.parent().unwrap()).unwrap();
std::fs::write(&fixture, serde_json::to_string_pretty(&actual).unwrap()).unwrap();
eprintln!("[snapshot] regenerated {}", fixture.display());
// Fail loudly so that accidentally setting KB_UPDATE_SNAPSHOTS
// in CI surfaces as a test failure rather than a silent
// overwrite + green run. Same fail-loud-instead-of-silent-pass
// philosophy as P3-2's `SNAPSHOT_HASH_BASELINE = 0` and P3-3's
// placeholder fixture guards.
panic!(
"[snapshot] regenerated {}, re-run without KB_UPDATE_SNAPSHOTS to verify pin",
fixture.display()
);
}
let expected: serde_json::Value =
serde_json::from_str(&std::fs::read_to_string(&fixture).unwrap_or_else(|_| {
panic!(
"missing snapshot fixture at {}; run with \
KB_UPDATE_SNAPSHOTS=1 to create",
fixture.display()
)
}))
.unwrap();
// Refuse to silently "pass" against the committed placeholder. The
// placeholder JSON carries a `_comment` field with regeneration
// instructions; production fixtures (a captured list) do not.
if expected.get("_comment").is_some() {
panic!(
"snapshot fixture is a placeholder — regenerate on AVX hardware then commit. \
Path: {}. To regenerate: \
`KB_UPDATE_SNAPSHOTS=1 cargo test -p kb-search -- --ignored hybrid_snapshot`.",
fixture.display()
);
}
assert_eq!(
actual, expected,
"hybrid snapshot drift; rerun with KB_UPDATE_SNAPSHOTS=1 to regenerate"
);
// Independent guard: fusion scores must be non-increasing across
// the result list (rrf is rank-biased, so this is the
// semantically-correct ordering invariant).
for w in hits.windows(2) {
assert!(
w[0].retrieval.fusion_score >= w[1].retrieval.fusion_score,
"fusion scores not in descending order: {} then {}",
w[0].retrieval.fusion_score,
w[1].retrieval.fusion_score
);
}
}

666
crates/kebab-search/tests/lexical.rs Normal file
View File

@@ -0,0 +1,666 @@
//! P2-2 integration tests for `LexicalRetriever`.
//!
//! Strategy: seed the SQLite store via raw inserts with `foreign_keys =
//! OFF` (mirroring the P2-1 FTS tests). This avoids dragging
//! `kb-parse-md` / `kb-normalize` / `kb-chunk` into kb-search's dev-deps,
//! which would violate the task's "Allowed deps" list.
use std::sync::Arc;
use kebab_config::Config;
use kebab_core::{IndexVersion, Lang, Retriever, SearchFilters, SearchMode, SearchQuery, TrustLevel};
use kebab_search::LexicalRetriever;
use kebab_store_sqlite::SqliteStore;
use rusqlite::Connection;
use tempfile::TempDir;
// ── Test scaffolding ─────────────────────────────────────────────────────
struct Env {
_temp: TempDir,
store: Arc<SqliteStore>,
db_path: std::path::PathBuf,
}
impl Env {
fn new() -> Self {
let temp = tempfile::tempdir().expect("tempdir");
let mut config = Config::defaults();
config.storage.data_dir = temp.path().to_string_lossy().into_owned();
let store = SqliteStore::open(&config).expect("open store");
store.run_migrations().expect("run migrations");
let db_path = temp.path().join("kb.sqlite");
Self {
_temp: temp,
store: Arc::new(store),
db_path,
}
}
/// Side-channel raw connection with FK enforcement off — same
/// trick used by P2-1's FTS tests so we can seed `chunks` /
/// `documents` directly without the full ingest graph.
fn raw_conn(&self) -> Connection {
let conn = Connection::open(&self.db_path).expect("open side conn");
conn.pragma_update(None, "foreign_keys", "OFF").unwrap();
conn
}
fn retriever(&self) -> LexicalRetriever {
LexicalRetriever::new(
Arc::clone(&self.store),
IndexVersion("v1.0".to_string()),
)
}
fn retriever_with_snippet_chars(&self, snippet_chars: usize) -> LexicalRetriever {
LexicalRetriever::with_settings(
Arc::clone(&self.store),
IndexVersion("v1.0".to_string()),
snippet_chars,
)
}
}
/// Minimal documents row. Many columns are NOT NULL and we don't care
/// about their exact values for retrieval tests, so we wedge in
/// reasonable defaults.
#[allow(clippy::too_many_arguments)]
fn insert_document(
conn: &Connection,
doc_id: &str,
workspace_path: &str,
title: &str,
lang: &str,
trust_level: &str,
tags: &[&str],
) {
// assets row first — documents.asset_id has a FK with ON DELETE
// RESTRICT but FKs are OFF on this connection. Still we insert a
// matching row so JOINs pick it up.
let asset_id = format!("{:0>32}", &doc_id[..1.min(doc_id.len())]); // 32-hex-ish
let asset_id = format!("{:0>32}", asset_id.chars().take(32).collect::<String>());
conn.execute(
"INSERT OR IGNORE INTO assets (
asset_id, source_uri, workspace_path, media_type, byte_len,
checksum, storage_kind, storage_path, discovered_at
) VALUES (?, 'file:///x', ?, '\"markdown\"', 0,
'd0', 'reference', '/x', '2024-01-01T00:00:00Z')",
rusqlite::params![asset_id, workspace_path],
)
.expect("insert asset");
conn.execute(
"INSERT INTO documents (
doc_id, asset_id, workspace_path, title, lang,
source_type, trust_level, parser_version,
doc_version, schema_version, metadata_json,
provenance_json, created_at, updated_at
) VALUES (?, ?, ?, ?, ?, 'markdown', ?, 'pv1', 1, 1,
'{}', '{\"events\":[]}',
'2024-01-01T00:00:00Z', '2024-01-01T00:00:00Z')",
rusqlite::params![doc_id, asset_id, workspace_path, title, lang, trust_level],
)
.expect("insert document");
for tag in tags {
conn.execute(
"INSERT INTO document_tags (doc_id, tag) VALUES (?, ?)",
rusqlite::params![doc_id, tag],
)
.expect("insert tag");
}
}
#[allow(clippy::too_many_arguments)]
fn insert_chunk(
conn: &Connection,
chunk_id: &str,
doc_id: &str,
text: &str,
heading_path: &[&str],
section_label: Option<&str>,
source_spans_json: &str,
chunker_version: &str,
) {
let heading_json = serde_json::to_string(heading_path).unwrap();
conn.execute(
"INSERT INTO chunks (
chunk_id, doc_id, text, heading_path_json, section_label,
source_spans_json, token_estimate, chunker_version,
policy_hash, block_ids_json, created_at
) VALUES (?, ?, ?, ?, ?, ?, 0, ?, 'h', '[]', '2024-01-01T00:00:00Z')",
rusqlite::params![
chunk_id,
doc_id,
text,
heading_json,
section_label,
source_spans_json,
chunker_version,
],
)
.expect("insert chunk");
}
/// Pad a short ID to the 32-hex shape kebab_core newtypes expect.
fn id32(prefix: &str) -> String {
let mut s = prefix.to_string();
while s.len() < 32 {
s.push('0');
}
s.truncate(32);
s
}
// ── Tests ────────────────────────────────────────────────────────────────
#[test]
fn lexical_empty_corpus_returns_empty_vec() {
let env = Env::new();
let r = env.retriever();
let q = SearchQuery {
text: "rust".to_string(),
mode: SearchMode::Lexical,
k: 10,
filters: SearchFilters::default(),
};
let hits = r.search(&q).expect("search");
assert!(hits.is_empty(), "empty corpus must yield empty Vec");
}
#[test]
fn lexical_empty_query_returns_empty_vec_without_db_hit() {
// Even with rows in the DB, a blank query must short-circuit to [].
let env = Env::new();
let conn = env.raw_conn();
insert_document(&conn, &id32("d"), "notes/a.md", "A", "en", "primary", &[]);
insert_chunk(
&conn,
&id32("c1"),
&id32("d"),
"rust cargo macros",
&["A"],
None,
r#"[{"kind":"line","start":1,"end":3}]"#,
"v1",
);
drop(conn);
let r = env.retriever();
for empty in ["", " ", "''"] {
let q = SearchQuery {
text: empty.to_string(),
mode: SearchMode::Lexical,
k: 5,
filters: SearchFilters::default(),
};
let hits = r.search(&q).unwrap();
assert!(hits.is_empty(), "query {empty:?} must yield empty Vec");
}
}
#[test]
fn lexical_single_doc_match_returns_one_hit_with_citation_round_trip() {
let env = Env::new();
let conn = env.raw_conn();
insert_document(&conn, &id32("d"), "notes/rust.md", "Rust Notes", "en", "primary", &[]);
insert_chunk(
&conn,
&id32("c1"),
&id32("d"),
"Rust borrow checker enforces ownership.",
&["Notes"],
Some("Notes"),
r#"[{"kind":"line","start":4,"end":4}]"#,
"v1",
);
drop(conn);
let r = env.retriever();
let q = SearchQuery {
text: "borrow".to_string(),
mode: SearchMode::Lexical,
k: 10,
filters: SearchFilters::default(),
};
let hits = r.search(&q).expect("search");
assert_eq!(hits.len(), 1);
let h = &hits[0];
assert_eq!(h.rank, 1);
assert_eq!(h.doc_path.0, "notes/rust.md");
assert_eq!(h.heading_path, vec!["Notes".to_string()]);
assert_eq!(h.section_label.as_deref(), Some("Notes"));
assert_eq!(h.retrieval.method, SearchMode::Lexical);
assert_eq!(h.retrieval.lexical_rank, Some(1));
assert!(h.retrieval.vector_score.is_none());
// Citation round-trips through `to_uri`/`parse` (line variant).
let uri = h.citation.to_uri();
let parsed = kebab_core::Citation::parse(&uri).expect("parse uri");
// Reparsed citation has section=None (URI fragment doesn't carry it),
// so compare by `to_uri` equivalence rather than struct equality.
assert_eq!(parsed.to_uri(), uri);
// Sanity: this is a Line citation matching the seeded source span.
assert_eq!(uri, "notes/rust.md#L4");
}
#[test]
fn lexical_snippet_length_capped_at_snippet_chars() {
let env = Env::new();
let conn = env.raw_conn();
insert_document(
&conn,
&id32("d"),
"notes/long.md",
"Long",
"en",
"primary",
&[],
);
// A text long enough that FTS5 might return a snippet > 80 chars
// when given a high word budget. We instead set a tight cap below
// and rely on `trim_snippet` as the backstop.
let mut text = String::new();
for _ in 0..50 {
text.push_str("alpha beta gamma delta epsilon ");
}
insert_chunk(
&conn,
&id32("c1"),
&id32("d"),
&text,
&["Long"],
None,
r#"[{"kind":"line","start":1,"end":1}]"#,
"v1",
);
drop(conn);
// Set snippet_chars to a known bound; the retriever clamps + trims
// any snippet to fit.
let r = env.retriever_with_snippet_chars(80);
let hits = r
.search(&SearchQuery {
text: "alpha".to_string(),
mode: SearchMode::Lexical,
k: 1,
filters: SearchFilters::default(),
})
.unwrap();
assert_eq!(hits.len(), 1);
assert!(
hits[0].snippet.chars().count() <= 80,
"snippet must be ≤ snippet_chars; got {} chars: {:?}",
hits[0].snippet.chars().count(),
hits[0].snippet
);
}
#[test]
fn lexical_filter_tags_any_excludes_untagged_docs() {
let env = Env::new();
let conn = env.raw_conn();
insert_document(&conn, &id32("d1"), "notes/a.md", "A", "en", "primary", &["rust"]);
insert_document(&conn, &id32("d2"), "notes/b.md", "B", "en", "primary", &["python"]);
insert_chunk(
&conn,
&id32("c1"),
&id32("d1"),
"ownership and borrow checker",
&["A"],
None,
r#"[{"kind":"line","start":1,"end":1}]"#,
"v1",
);
insert_chunk(
&conn,
&id32("c2"),
&id32("d2"),
"borrow semantics in python",
&["B"],
None,
r#"[{"kind":"line","start":1,"end":1}]"#,
"v1",
);
drop(conn);
let r = env.retriever();
let q = SearchQuery {
text: "borrow".to_string(),
mode: SearchMode::Lexical,
k: 10,
filters: SearchFilters {
tags_any: vec!["rust".to_string()],
..Default::default()
},
};
let hits = r.search(&q).unwrap();
assert_eq!(hits.len(), 1, "tags_any=[rust] must exclude python doc");
assert_eq!(hits[0].doc_path.0, "notes/a.md");
}
#[test]
fn lexical_filter_lang_and_trust_min_compose() {
let env = Env::new();
let conn = env.raw_conn();
insert_document(&conn, &id32("d1"), "ko/a.md", "A", "ko", "primary", &[]);
insert_document(&conn, &id32("d2"), "en/b.md", "B", "en", "primary", &[]);
insert_document(&conn, &id32("d3"), "en/c.md", "C", "en", "generated", &[]);
for (cid, did, body) in [
("c1", "d1", "검색 키워드 alpha"),
("c2", "d2", "alpha bravo"),
("c3", "d3", "alpha gamma"),
] {
insert_chunk(
&conn,
&id32(cid),
&id32(did),
body,
&[],
None,
r#"[{"kind":"line","start":1,"end":1}]"#,
"v1",
);
}
drop(conn);
let r = env.retriever();
// lang=en + trust_min=secondary → only d2 (primary ≥ secondary).
let hits = r
.search(&SearchQuery {
text: "alpha".to_string(),
mode: SearchMode::Lexical,
k: 10,
filters: SearchFilters {
lang: Some(Lang("en".to_string())),
trust_min: Some(TrustLevel::Secondary),
..Default::default()
},
})
.unwrap();
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].doc_path.0, "en/b.md");
}
#[test]
fn lexical_filter_path_glob_does_not_cross_slash() {
let env = Env::new();
let conn = env.raw_conn();
insert_document(&conn, &id32("d1"), "notes/a.md", "A", "en", "primary", &[]);
insert_document(&conn, &id32("d2"), "notes/sub/b.md", "B", "en", "primary", &[]);
insert_chunk(
&conn,
&id32("c1"),
&id32("d1"),
"shared keyword",
&[],
None,
r#"[{"kind":"line","start":1,"end":1}]"#,
"v1",
);
insert_chunk(
&conn,
&id32("c2"),
&id32("d2"),
"shared keyword",
&[],
None,
r#"[{"kind":"line","start":1,"end":1}]"#,
"v1",
);
drop(conn);
let r = env.retriever();
let hits = r
.search(&SearchQuery {
text: "keyword".to_string(),
mode: SearchMode::Lexical,
k: 10,
filters: SearchFilters {
path_glob: Some("notes/*.md".to_string()),
..Default::default()
},
})
.unwrap();
let paths: Vec<&str> = hits.iter().map(|h| h.doc_path.0.as_str()).collect();
assert_eq!(paths, vec!["notes/a.md"], "* must not match across `/`");
}
#[test]
fn lexical_citation_round_trip_against_first_source_span() {
let env = Env::new();
let conn = env.raw_conn();
insert_document(&conn, &id32("d"), "notes/m.md", "M", "en", "primary", &[]);
insert_chunk(
&conn,
&id32("c1"),
&id32("d"),
"echo bravo",
&[],
None,
// Two spans; the citation uses the first.
r#"[{"kind":"line","start":12,"end":34},{"kind":"line","start":60,"end":61}]"#,
"v1",
);
drop(conn);
let r = env.retriever();
let hits = r
.search(&SearchQuery {
text: "bravo".to_string(),
mode: SearchMode::Lexical,
k: 1,
filters: SearchFilters::default(),
})
.unwrap();
assert_eq!(hits.len(), 1);
let uri = hits[0].citation.to_uri();
assert_eq!(uri, "notes/m.md#L12-L34");
let parsed = kebab_core::Citation::parse(&uri).unwrap();
assert_eq!(parsed.to_uri(), uri);
}
#[test]
fn lexical_top_score_within_unit_interval_three_chunks() {
let env = Env::new();
let conn = env.raw_conn();
insert_document(&conn, &id32("d"), "notes/r.md", "R", "en", "primary", &[]);
// Three chunks of varying relevance to the query 'alpha':
// c1: alpha alpha alpha (best)
// c2: alpha bravo
// c3: bravo charlie alpha (one occurrence)
for (cid, body) in [
("c1", "alpha alpha alpha keyword"),
("c2", "alpha bravo charlie"),
("c3", "bravo charlie alpha"),
] {
insert_chunk(
&conn,
&id32(cid),
&id32("d"),
body,
&[],
None,
r#"[{"kind":"line","start":1,"end":1}]"#,
"v1",
);
}
drop(conn);
let r = env.retriever();
let hits = r
.search(&SearchQuery {
text: "alpha".to_string(),
mode: SearchMode::Lexical,
k: 10,
filters: SearchFilters::default(),
})
.unwrap();
assert!(!hits.is_empty(), "must surface at least one hit");
let top = hits[0].retrieval.fusion_score;
assert!(
top > 0.0 && top <= 1.0,
"top normalized score must be in (0, 1]; got {top}"
);
// All scores in [0, 1].
for h in &hits {
let s = h.retrieval.fusion_score;
assert!((0.0..=1.0).contains(&s), "hit score {s} out of [0, 1]");
// lexical_score and fusion_score equal in lexical-only mode.
assert_eq!(h.retrieval.lexical_score, Some(s));
}
// bm25 should rank c1 (3 occurrences) above c2 / c3.
assert!(hits[0].chunk_id.0.starts_with("c1"));
}
#[test]
fn lexical_determinism_same_query_twice() {
let env = Env::new();
let conn = env.raw_conn();
insert_document(&conn, &id32("d"), "notes/r.md", "R", "en", "primary", &[]);
for (cid, body) in [
("c1", "alpha alpha"),
("c2", "alpha bravo"),
("c3", "alpha charlie"),
("c4", "alpha delta"),
] {
insert_chunk(
&conn,
&id32(cid),
&id32("d"),
body,
&[],
None,
r#"[{"kind":"line","start":1,"end":1}]"#,
"v1",
);
}
drop(conn);
let r = env.retriever();
let q = SearchQuery {
text: "alpha".to_string(),
mode: SearchMode::Lexical,
k: 10,
filters: SearchFilters::default(),
};
let a = r.search(&q).unwrap();
let b = r.search(&q).unwrap();
assert_eq!(a, b, "same DB + same query must yield identical Vec<SearchHit>");
}
#[test]
fn lexical_determinism_chunk_id_tiebreaker_on_equal_bm25() {
// Two chunks with byte-identical text + length → identical bm25 scores
// for any `MATCH` against them. The retriever must fall back to
// `chunk_id` ordering so the result is stable across runs.
let env = Env::new();
let conn = env.raw_conn();
insert_document(&conn, &id32("d"), "notes/tie.md", "Tie", "en", "primary", &[]);
let cid_a = id32("aaaa");
let cid_b = id32("bbbb");
assert!(cid_a < cid_b, "test premise: aaaa-id sorts before bbbb-id");
for cid in [&cid_a, &cid_b] {
insert_chunk(
&conn,
cid,
&id32("d"),
"alpha bravo charlie",
&[],
None,
r#"[{"kind":"line","start":1,"end":1}]"#,
"v1",
);
}
drop(conn);
let r = env.retriever();
let q = SearchQuery {
text: "alpha".to_string(),
mode: SearchMode::Lexical,
k: 10,
filters: SearchFilters::default(),
};
let a = r.search(&q).unwrap();
let b = r.search(&q).unwrap();
assert_eq!(a.len(), 2, "both chunks should match");
// bm25 must be equal for byte-identical chunks; the secondary sort
// by chunk_id pins the order.
assert!(
(a[0].retrieval.fusion_score - a[1].retrieval.fusion_score).abs() < 1e-9,
"byte-identical chunks must score equally; got {} vs {}",
a[0].retrieval.fusion_score,
a[1].retrieval.fusion_score
);
assert!(
a[0].chunk_id.0 < a[1].chunk_id.0,
"tiebreaker must order by chunk_id ascending; got {} then {}",
a[0].chunk_id.0,
a[1].chunk_id.0
);
assert_eq!(a, b, "tiebreaker order must be stable across runs");
}
#[test]
fn lexical_index_version_is_returned_unchanged() {
let env = Env::new();
let r = LexicalRetriever::new(
Arc::clone(&env.store),
IndexVersion("custom-label-1".to_string()),
);
assert_eq!(r.index_version().0, "custom-label-1");
}
#[test]
fn lexical_snapshot_run_1() {
// Pinned snapshot. A small, deterministic corpus; the JSON shape of
// `Vec<SearchHit>` for a fixed query is checked verbatim against
// `tests/fixtures/search/lexical/run-1.json`. Update both sides in
// the same commit when intentional changes ship.
// Stable because rusqlite ships bundled SQLite — a tokenizer/bm25 algorithm change in a future SQLite bump will require regenerating run-1.json via `KB_UPDATE_SNAPSHOTS=1`.
let env = Env::new();
let conn = env.raw_conn();
insert_document(&conn, &id32("d"), "notes/snap.md", "Snap", "en", "primary", &[]);
for (cid, body, span) in [
(
"c1",
"alpha bravo charlie",
r#"[{"kind":"line","start":1,"end":2}]"#,
),
(
"c2",
"bravo only here",
r#"[{"kind":"line","start":4,"end":5}]"#,
),
(
"c3",
"alpha alpha",
r#"[{"kind":"line","start":7,"end":8}]"#,
),
] {
insert_chunk(&conn, &id32(cid), &id32("d"), body, &["Snap"], Some("Snap"), span, "v1");
}
drop(conn);
let r = env.retriever();
let hits = r
.search(&SearchQuery {
text: "alpha".to_string(),
mode: SearchMode::Lexical,
k: 10,
filters: SearchFilters::default(),
})
.unwrap();
let actual = serde_json::to_value(&hits).unwrap();
let baseline_path =
std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/search/lexical/run-1.json");
if std::env::var_os("KB_UPDATE_SNAPSHOTS").is_some() {
std::fs::write(&baseline_path, serde_json::to_string_pretty(&actual).unwrap()).unwrap();
}
let baseline_text = std::fs::read_to_string(&baseline_path)
.expect("baseline snapshot must exist; run with KB_UPDATE_SNAPSHOTS=1 to seed");
let expected: serde_json::Value = serde_json::from_str(&baseline_text).unwrap();
assert_eq!(actual, expected, "lexical run-1 snapshot drift");
}