llm_kernel/search/
provider.rs

1//! Pluggable search backends behind a unified sync interface.
2//!
3//! [`SearchProvider`] is the common trait every ranking backend implements.
4//! [`KeywordIndex`] is a minimal, dependency-free term-frequency backend used
5//! for tests and simple keyword search.
6
7use crate::search::SearchResult;
8
9/// Unified sync interface for search backends producing ranked results.
10pub trait SearchProvider: Send + Sync {
11    /// Human-readable backend name.
12    fn name(&self) -> &'static str;
13    /// Run a query, returning up to `limit` ranked results.
14    fn search(&self, query: &str, limit: usize) -> crate::error::Result<Vec<SearchResult>>;
15}
16
17/// A simple keyword (term-frequency) index over a fixed set of documents.
18///
19/// Scores each document by the total count of query-term occurrences in its
20/// text (case-insensitive substring matching — a query term `rust` will match
21/// inside `trust` or `rustic`). No stemming, no IDF weighting — a building
22/// block, not a BM25 replacement.
23pub struct KeywordIndex {
24    /// Indexed documents as `(id, text)` pairs.
25    docs: Vec<(String, String)>,
26}
27
28impl KeywordIndex {
29    /// Build an index from `(id, text)` document pairs.
30    pub fn new(docs: Vec<(String, String)>) -> Self {
31        Self { docs }
32    }
33}
34
35impl SearchProvider for KeywordIndex {
36    fn name(&self) -> &'static str {
37        "keyword"
38    }
39
40    fn search(&self, query: &str, limit: usize) -> crate::error::Result<Vec<SearchResult>> {
41        if query.is_empty() {
42            return Ok(Vec::new());
43        }
44        let query_lower = query.to_lowercase();
45        let terms: Vec<&str> = query_lower.split_ascii_whitespace().collect();
46
47        let mut scored: Vec<SearchResult> = self
48            .docs
49            .iter()
50            .filter_map(|(id, text)| {
51                let text_lower = text.to_lowercase();
52                let mut count: usize = 0;
53                for term in &terms {
54                    count += text_lower.matches(term).count();
55                }
56                if count == 0 {
57                    None
58                } else {
59                    Some(SearchResult {
60                        id: id.clone(),
61                        score: count as f32,
62                        text: text.clone(),
63                    })
64                }
65            })
66            .collect();
67
68        scored.sort_by(|a, b| {
69            b.score
70                .partial_cmp(&a.score)
71                .unwrap_or(std::cmp::Ordering::Equal)
72        });
73        scored.truncate(limit);
74        Ok(scored)
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81
82    fn sample_docs() -> Vec<(String, String)> {
83        vec![
84            (
85                "d1".to_string(),
86                "the rust programming language is fast".to_string(),
87            ),
88            (
89                "d2".to_string(),
90                "python is a popular programming language".to_string(),
91            ),
92            (
93                "d3".to_string(),
94                "rust has zero cost abstractions".to_string(),
95            ),
96            ("d4".to_string(), "cooking recipes for dinner".to_string()),
97        ]
98    }
99
100    #[test]
101    fn all_terms_outranks_fewer() {
102        let index = KeywordIndex::new(sample_docs());
103        let results = index.search("rust programming language", 10).unwrap();
104        // d1 contains all three query terms; d2 contains two; d3 contains one.
105        assert_eq!(results[0].id, "d1");
106        let d1 = results.iter().find(|r| r.id == "d1").unwrap().score;
107        let d2 = results.iter().find(|r| r.id == "d2").unwrap().score;
108        let d3 = results.iter().find(|r| r.id == "d3").unwrap().score;
109        assert!(d1 > d2);
110        assert!(d2 > d3);
111    }
112
113    #[test]
114    fn empty_query_returns_empty() {
115        let index = KeywordIndex::new(sample_docs());
116        let results = index.search("", 10).unwrap();
117        assert!(results.is_empty());
118    }
119
120    #[test]
121    fn limit_is_respected() {
122        let index = KeywordIndex::new(sample_docs());
123        let results = index.search("programming language", 1).unwrap();
124        assert_eq!(results.len(), 1);
125    }
126
127    #[test]
128    fn unknown_term_returns_only_matches() {
129        let index = KeywordIndex::new(sample_docs());
130        // "rust" appears only in d1 and d3; "xyzzy" appears nowhere.
131        let results = index.search("rust xyzzy", 10).unwrap();
132        let ids: Vec<&str> = results.iter().map(|r| r.id.as_str()).collect();
133        assert!(ids.contains(&"d1"));
134        assert!(ids.contains(&"d3"));
135        assert!(!ids.contains(&"d2"));
136        assert!(!ids.contains(&"d4"));
137    }
138
139    #[test]
140    fn name_is_keyword() {
141        let index = KeywordIndex::new(vec![]);
142        assert_eq!(index.name(), "keyword");
143    }
144
145    #[test]
146    fn zero_score_docs_excluded() {
147        let index = KeywordIndex::new(sample_docs());
148        let results = index.search("rust", 10).unwrap();
149        // d4 does not contain "rust" -> excluded entirely.
150        assert!(results.iter().all(|r| r.id != "d4"));
151    }
152}
llm_kernel/search/provider.rs

llm_kernel/search/
provider.rs
