Skip to main content

scour/
lib.rs

1//! # Scour
2//!
3//! Embeddable **hybrid search primitives** for Rust — the in-process core
4//! of a retrieval pipeline, with **zero dependencies**:
5//!
6//! - [`bm25::Bm25Index`] — Okapi BM25 keyword search (inverted index,
7//!   Porter stemming, stopword removal)
8//! - [`hnsw::HnswIndex`] — HNSW approximate nearest neighbor vector index
9//!   (cosine distance, soft deletes, deterministic construction)
10//! - [`fuse::rrf_fuse`] — Reciprocal Rank Fusion to merge lexical and
11//!   semantic rankings
12//! - [`chunk::chunk_text`] — UTF-8-safe, boundary-aware text chunking for
13//!   embedding pipelines
14//! - [`HybridIndex`] — the four combined: one type that indexes text +
15//!   vectors and serves fused hybrid queries
16//!
17//! ## Example
18//!
19//! ```
20//! use scour::HybridIndex;
21//!
22//! let mut index = HybridIndex::new(3); // 3-dim embeddings for the demo
23//! index.add("doc-a", "rust systems programming", &[1.0, 0.0, 0.0]);
24//! index.add("doc-b", "gardening in spring", &[0.0, 1.0, 0.0]);
25//!
26//! // Hybrid query: keyword text + query embedding, fused with RRF.
27//! let results = index.search("rust programming", &[0.9, 0.1, 0.0], 2);
28//! assert_eq!(results[0].0, "doc-a");
29//! ```
30//!
31//! Bring your own embeddings: Scour is deliberately model-agnostic. Any
32//! `&[f32]` works — ONNX, candle, an HTTP embedding service, or test
33//! fixtures.
34
35pub mod bm25;
36pub mod chunk;
37pub mod corpus;
38pub mod embed;
39pub mod fuse;
40pub mod hnsw;
41pub mod text;
42
43pub use bm25::Bm25Index;
44pub use chunk::{chunk_text, chunk_text_with_overlap};
45pub use embed::{embed, DEMO_DIM};
46pub use fuse::{rrf_fuse, rrf_fuse_scored, DEFAULT_RRF_K};
47pub use hnsw::{cosine_distance, HnswIndex, HnswParams};
48
49/// A combined lexical + vector index serving RRF-fused hybrid queries.
50///
51/// Wraps a [`Bm25Index`] and an [`HnswIndex`] under one id space.
52pub struct HybridIndex {
53    lexical: Bm25Index,
54    vector: HnswIndex,
55    rrf_k: f64,
56}
57
58impl HybridIndex {
59    /// Create a hybrid index for embeddings of `dimensions`.
60    pub fn new(dimensions: usize) -> Self {
61        Self::with_params(dimensions, HnswParams::default(), DEFAULT_RRF_K)
62    }
63
64    pub fn with_params(dimensions: usize, hnsw: HnswParams, rrf_k: f64) -> Self {
65        Self {
66            lexical: Bm25Index::new(),
67            vector: HnswIndex::with_params(dimensions, hnsw),
68            rrf_k,
69        }
70    }
71
72    /// Index a document under `id` with its text and embedding.
73    /// Re-adding an id replaces both representations.
74    pub fn add(&mut self, id: &str, text: &str, embedding: &[f32]) {
75        self.lexical.add_document(id, text);
76        self.vector.insert(id, embedding);
77    }
78
79    /// Remove a document from both indexes. Returns whether the id existed
80    /// in the vector index.
81    pub fn remove(&mut self, id: &str) -> bool {
82        self.lexical.remove_document(id);
83        self.vector.remove(id)
84    }
85
86    /// Number of live documents (vector-index count).
87    pub fn len(&self) -> usize {
88        self.vector.len()
89    }
90
91    pub fn is_empty(&self) -> bool {
92        self.vector.is_empty()
93    }
94
95    /// Hybrid query: BM25 over `query_text`, ANN over `query_embedding`,
96    /// fused with RRF. Returns `(id, fused_score)` best-first.
97    ///
98    /// Each leg retrieves `k * 3` candidates (a standard over-fetch so the
99    /// fusion has material to work with), and the fused list is truncated
100    /// to `k`.
101    pub fn search(
102        &self,
103        query_text: &str,
104        query_embedding: &[f32],
105        k: usize,
106    ) -> Vec<(String, f64)> {
107        if k == 0 {
108            return Vec::new();
109        }
110        let fetch = k.saturating_mul(3);
111        let lexical = self.lexical.search(query_text, fetch);
112        let semantic = self.vector.search(query_embedding, fetch);
113
114        let mut fused = rrf_fuse_scored(&lexical, &semantic, self.rrf_k);
115        fused.truncate(k);
116        fused
117    }
118
119    /// Lexical-only query (BM25).
120    pub fn search_lexical(&self, query: &str, k: usize) -> Vec<(String, f64)> {
121        self.lexical.search(query, k)
122    }
123
124    /// Semantic-only query (HNSW, cosine distance — smaller is closer).
125    pub fn search_semantic(&self, query: &[f32], k: usize) -> Vec<(String, f32)> {
126        self.vector.search(query, k)
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133
134    #[test]
135    fn hybrid_prefers_doc_matching_both_signals() {
136        let mut index = HybridIndex::new(4);
137        // doc-a: strong on both keyword and vector.
138        index.add(
139            "doc-a",
140            "rust async runtime internals",
141            &[1.0, 0.0, 0.0, 0.0],
142        );
143        // doc-b: keyword match only.
144        index.add("doc-b", "rust cookbook recipes", &[0.0, 1.0, 0.0, 0.0]);
145        // doc-c: vector match only.
146        index.add("doc-c", "tokio scheduler design", &[0.9, 0.1, 0.0, 0.0]);
147
148        let results = index.search("rust runtime", &[1.0, 0.0, 0.0, 0.0], 3);
149        assert_eq!(results[0].0, "doc-a");
150    }
151
152    #[test]
153    fn remove_drops_from_both_legs() {
154        let mut index = HybridIndex::new(2);
155        index.add("x", "unique pelican words", &[1.0, 0.0]);
156        index.add("y", "other content", &[0.0, 1.0]);
157
158        assert!(index.remove("x"));
159        let results = index.search("pelican", &[1.0, 0.0], 5);
160        assert!(results.iter().all(|(id, _)| id != "x"));
161        assert_eq!(index.len(), 1);
162    }
163
164    #[test]
165    fn k_zero_returns_empty() {
166        let mut index = HybridIndex::new(2);
167        index.add("x", "abc", &[1.0, 0.0]);
168        assert!(index.search("abc", &[1.0, 0.0], 0).is_empty());
169    }
170}