semtree_rag/
filter.rs

1use semtree_core::{Chunk, ChunkKind, Language};
2
3use crate::RagError;
4
5/// How many candidates to over-fetch per requested result when a filter is
6/// active, and the floor below which over-fetching is not worth tuning.
7const OVERFETCH_FACTOR: usize = 10;
8const OVERFETCH_FLOOR: usize = 50;
9
10/// Metadata narrowing applied to search hits after ranking.
11///
12/// Neither the vector index nor BM25 knows about languages or chunk kinds, so
13/// narrowing happens here, against the chunk each hit resolves to. Because it
14/// runs *after* ranking, the ranker has to be asked for more candidates than the
15/// caller wants - see [`fetch_size`](SearchFilters::fetch_size).
16///
17/// ```
18/// use semtree_rag::SearchFilters;
19///
20/// let filters = SearchFilters::default()
21///     .with_language_names(["rust"])?
22///     .with_kind_names(["fn"])?
23///     .with_path("src/");
24/// # Ok::<(), semtree_rag::RagError>(())
25/// ```
26#[derive(Debug, Clone, Default)]
27pub struct SearchFilters {
28    languages: Vec<Language>,
29    kinds: Vec<ChunkKind>,
30    path: Option<String>,
31}
32
33impl SearchFilters {
34    /// Keep only chunks in these languages. Empty means "any".
35    pub fn with_languages(mut self, languages: impl IntoIterator<Item = Language>) -> Self {
36        self.languages.extend(languages);
37        self
38    }
39
40    /// Keep only chunks of these kinds. Empty means "any".
41    pub fn with_kinds(mut self, kinds: impl IntoIterator<Item = ChunkKind>) -> Self {
42        self.kinds.extend(kinds);
43        self
44    }
45
46    /// Keep only chunks whose path contains `needle`.
47    pub fn with_path(mut self, needle: impl Into<String>) -> Self {
48        self.path = Some(needle.into());
49        self
50    }
51
52    /// Like [`with_languages`](Self::with_languages), from user-supplied names.
53    /// An unrecognized name is an error rather than a filter that matches
54    /// nothing, so a typo surfaces instead of looking like an empty index.
55    pub fn with_language_names<S: AsRef<str>>(
56        self,
57        names: impl IntoIterator<Item = S>,
58    ) -> Result<Self, RagError> {
59        let languages = names
60            .into_iter()
61            .map(|name| {
62                Language::from_name(name.as_ref()).ok_or_else(|| {
63                    RagError::Filter(format!(
64                        "unknown language '{}' (expected one of: {})",
65                        name.as_ref(),
66                        join_names(Language::ALL)
67                    ))
68                })
69            })
70            .collect::<Result<Vec<_>, _>>()?;
71        Ok(self.with_languages(languages))
72    }
73
74    /// Like [`with_kinds`](Self::with_kinds), from user-supplied names.
75    pub fn with_kind_names<S: AsRef<str>>(
76        self,
77        names: impl IntoIterator<Item = S>,
78    ) -> Result<Self, RagError> {
79        let kinds = names
80            .into_iter()
81            .map(|name| {
82                ChunkKind::from_name(name.as_ref()).ok_or_else(|| {
83                    RagError::Filter(format!(
84                        "unknown kind '{}' (expected one of: {})",
85                        name.as_ref(),
86                        join_names(ChunkKind::ALL)
87                    ))
88                })
89            })
90            .collect::<Result<Vec<_>, _>>()?;
91        Ok(self.with_kinds(kinds))
92    }
93
94    /// Whether every hit passes untouched.
95    pub fn is_empty(&self) -> bool {
96        self.languages.is_empty() && self.kinds.is_empty() && self.path.is_none()
97    }
98
99    /// Whether `chunk` survives the filters.
100    pub fn matches(&self, chunk: &Chunk) -> bool {
101        if !self.languages.is_empty() && !self.languages.contains(&chunk.language) {
102            return false;
103        }
104        if !self.kinds.is_empty() && !self.kinds.contains(&chunk.kind) {
105            return false;
106        }
107        if let Some(needle) = &self.path
108            && !chunk.path.to_string_lossy().contains(needle.as_str())
109        {
110            return false;
111        }
112        true
113    }
114
115    /// How many candidates the ranker should return so that `top_k` survive the
116    /// filters. Unfiltered searches ask for exactly what they need.
117    pub fn fetch_size(&self, top_k: usize) -> usize {
118        if self.is_empty() {
119            top_k
120        } else {
121            (top_k * OVERFETCH_FACTOR).max(OVERFETCH_FLOOR)
122        }
123    }
124}
125
126fn join_names<T: std::fmt::Display>(values: &[T]) -> String {
127    values
128        .iter()
129        .map(|v| v.to_string())
130        .collect::<Vec<_>>()
131        .join(", ")
132}
133
134#[cfg(test)]
135mod tests {
136    use semtree_core::Span;
137
138    use super::*;
139
140    fn chunk(id: &str, language: Language, kind: ChunkKind, path: &str) -> Chunk {
141        Chunk {
142            id: id.to_string(),
143            path: path.into(),
144            language,
145            kind,
146            name: Some(id.to_string()),
147            content: String::new(),
148            span: Span::new(0, 0, 0, 0),
149            doc: None,
150        }
151    }
152
153    #[test]
154    fn empty_filters_match_everything() {
155        let filters = SearchFilters::default();
156        assert!(filters.is_empty());
157        assert!(filters.matches(&chunk("a", Language::Zig, ChunkKind::Struct, "src/a.zig")));
158    }
159
160    #[test]
161    fn filters_combine_as_and() {
162        let filters = SearchFilters::default()
163            .with_languages([Language::Rust])
164            .with_kinds([ChunkKind::Function])
165            .with_path("src/");
166
167        assert!(filters.matches(&chunk(
168            "ok",
169            Language::Rust,
170            ChunkKind::Function,
171            "src/a.rs"
172        )));
173        // Right language and kind, wrong path.
174        assert!(!filters.matches(&chunk(
175            "no",
176            Language::Rust,
177            ChunkKind::Function,
178            "tests/a.rs"
179        )));
180        // Right path and kind, wrong language.
181        assert!(!filters.matches(&chunk("no", Language::Go, ChunkKind::Function, "src/a.go")));
182    }
183
184    #[test]
185    fn unknown_names_are_rejected_rather_than_matching_nothing() {
186        let err = SearchFilters::default()
187            .with_language_names(["rust", "cobol"])
188            .unwrap_err();
189        assert!(err.to_string().contains("cobol"), "names the bad value");
190        assert!(err.to_string().contains("rust"), "lists valid values");
191
192        assert!(SearchFilters::default().with_kind_names(["gizmo"]).is_err());
193    }
194
195    #[test]
196    fn every_supported_language_is_accepted_by_name() {
197        // A language the parser handles but the filter rejects would be
198        // invisible from the CLI and the MCP server alike.
199        for lang in Language::ALL {
200            assert!(
201                SearchFilters::default()
202                    .with_language_names([lang.to_string()])
203                    .is_ok(),
204                "{lang} is not accepted as a filter"
205            );
206        }
207    }
208
209    #[test]
210    fn overfetches_only_when_filtering() {
211        assert_eq!(SearchFilters::default().fetch_size(5), 5);
212        let filtered = SearchFilters::default().with_path("src/");
213        assert!(filtered.fetch_size(5) >= OVERFETCH_FLOOR);
214        assert!(filtered.fetch_size(100) >= 100 * OVERFETCH_FACTOR);
215    }
216}
semtree_rag/filter.rs

semtree_rag/
filter.rs
