semtree-rag 0.5.0

use semtree_core::{Chunk, ChunkKind, Language};

use crate::RagError;

/// How many candidates to over-fetch per requested result when a filter is
/// active, and the floor below which over-fetching is not worth tuning.
const OVERFETCH_FACTOR: usize = 10;
const OVERFETCH_FLOOR: usize = 50;

/// Metadata narrowing applied to search hits after ranking.
///
/// Neither the vector index nor BM25 knows about languages or chunk kinds, so
/// narrowing happens here, against the chunk each hit resolves to. Because it
/// runs *after* ranking, the ranker has to be asked for more candidates than the
/// caller wants - see [`fetch_size`](SearchFilters::fetch_size).
///
/// ```
/// use semtree_rag::SearchFilters;
///
/// let filters = SearchFilters::default()
///     .with_language_names(["rust"])?
///     .with_kind_names(["fn"])?
///     .with_path("src/");
/// # Ok::<(), semtree_rag::RagError>(())
/// ```
#[derive(Debug, Clone, Default)]
pub struct SearchFilters {
    languages: Vec<Language>,
    kinds: Vec<ChunkKind>,
    path: Option<String>,
}

impl SearchFilters {
    /// Keep only chunks in these languages. Empty means "any".
    pub fn with_languages(mut self, languages: impl IntoIterator<Item = Language>) -> Self {
        self.languages.extend(languages);
        self
    }

    /// Keep only chunks of these kinds. Empty means "any".
    pub fn with_kinds(mut self, kinds: impl IntoIterator<Item = ChunkKind>) -> Self {
        self.kinds.extend(kinds);
        self
    }

    /// Keep only chunks whose path contains `needle`.
    pub fn with_path(mut self, needle: impl Into<String>) -> Self {
        self.path = Some(needle.into());
        self
    }

    /// Like [`with_languages`](Self::with_languages), from user-supplied names.
    /// An unrecognized name is an error rather than a filter that matches
    /// nothing, so a typo surfaces instead of looking like an empty index.
    pub fn with_language_names<S: AsRef<str>>(
        self,
        names: impl IntoIterator<Item = S>,
    ) -> Result<Self, RagError> {
        let languages = names
            .into_iter()
            .map(|name| {
                Language::from_name(name.as_ref()).ok_or_else(|| {
                    RagError::Filter(format!(
                        "unknown language '{}' (expected one of: {})",
                        name.as_ref(),
                        join_names(Language::ALL)
                    ))
                })
            })
            .collect::<Result<Vec<_>, _>>()?;
        Ok(self.with_languages(languages))
    }

    /// Like [`with_kinds`](Self::with_kinds), from user-supplied names.
    pub fn with_kind_names<S: AsRef<str>>(
        self,
        names: impl IntoIterator<Item = S>,
    ) -> Result<Self, RagError> {
        let kinds = names
            .into_iter()
            .map(|name| {
                ChunkKind::from_name(name.as_ref()).ok_or_else(|| {
                    RagError::Filter(format!(
                        "unknown kind '{}' (expected one of: {})",
                        name.as_ref(),
                        join_names(ChunkKind::ALL)
                    ))
                })
            })
            .collect::<Result<Vec<_>, _>>()?;
        Ok(self.with_kinds(kinds))
    }

    /// Whether every hit passes untouched.
    pub fn is_empty(&self) -> bool {
        self.languages.is_empty() && self.kinds.is_empty() && self.path.is_none()
    }

    /// Whether `chunk` survives the filters.
    pub fn matches(&self, chunk: &Chunk) -> bool {
        if !self.languages.is_empty() && !self.languages.contains(&chunk.language) {
            return false;
        }
        if !self.kinds.is_empty() && !self.kinds.contains(&chunk.kind) {
            return false;
        }
        if let Some(needle) = &self.path
            && !chunk.path.to_string_lossy().contains(needle.as_str())
        {
            return false;
        }
        true
    }

    /// How many candidates the ranker should return so that `top_k` survive the
    /// filters. Unfiltered searches ask for exactly what they need.
    pub fn fetch_size(&self, top_k: usize) -> usize {
        if self.is_empty() {
            top_k
        } else {
            (top_k * OVERFETCH_FACTOR).max(OVERFETCH_FLOOR)
        }
    }
}

fn join_names<T: std::fmt::Display>(values: &[T]) -> String {
    values
        .iter()
        .map(|v| v.to_string())
        .collect::<Vec<_>>()
        .join(", ")
}

#[cfg(test)]
mod tests {
    use semtree_core::Span;

    use super::*;

    fn chunk(id: &str, language: Language, kind: ChunkKind, path: &str) -> Chunk {
        Chunk {
            id: id.to_string(),
            path: path.into(),
            language,
            kind,
            name: Some(id.to_string()),
            content: String::new(),
            span: Span::new(0, 0, 0, 0),
            doc: None,
        }
    }

    #[test]
    fn empty_filters_match_everything() {
        let filters = SearchFilters::default();
        assert!(filters.is_empty());
        assert!(filters.matches(&chunk("a", Language::Zig, ChunkKind::Struct, "src/a.zig")));
    }

    #[test]
    fn filters_combine_as_and() {
        let filters = SearchFilters::default()
            .with_languages([Language::Rust])
            .with_kinds([ChunkKind::Function])
            .with_path("src/");

        assert!(filters.matches(&chunk(
            "ok",
            Language::Rust,
            ChunkKind::Function,
            "src/a.rs"
        )));
        // Right language and kind, wrong path.
        assert!(!filters.matches(&chunk(
            "no",
            Language::Rust,
            ChunkKind::Function,
            "tests/a.rs"
        )));
        // Right path and kind, wrong language.
        assert!(!filters.matches(&chunk("no", Language::Go, ChunkKind::Function, "src/a.go")));
    }

    #[test]
    fn unknown_names_are_rejected_rather_than_matching_nothing() {
        let err = SearchFilters::default()
            .with_language_names(["rust", "cobol"])
            .unwrap_err();
        assert!(err.to_string().contains("cobol"), "names the bad value");
        assert!(err.to_string().contains("rust"), "lists valid values");

        assert!(SearchFilters::default().with_kind_names(["gizmo"]).is_err());
    }

    #[test]
    fn every_supported_language_is_accepted_by_name() {
        // A language the parser handles but the filter rejects would be
        // invisible from the CLI and the MCP server alike.
        for lang in Language::ALL {
            assert!(
                SearchFilters::default()
                    .with_language_names([lang.to_string()])
                    .is_ok(),
                "{lang} is not accepted as a filter"
            );
        }
    }

    #[test]
    fn overfetches_only_when_filtering() {
        assert_eq!(SearchFilters::default().fetch_size(5), 5);
        let filtered = SearchFilters::default().with_path("src/");
        assert!(filtered.fetch_size(5) >= OVERFETCH_FLOOR);
        assert!(filtered.fetch_size(100) >= 100 * OVERFETCH_FACTOR);
    }
}