llm-prefix-match 0.1.0

/*!
llm-prefix-match: gate LLM outputs based on expected prefix patterns.

When you need a model to start its response a certain way (e.g. JSON,
a keyword, or a specific token), this crate lets you check, assert, or
strip that prefix before passing the response downstream.

```rust
use llm_prefix_match::{PrefixMatcher, MatchResult};

let m = PrefixMatcher::new().require_any(&["YES", "NO"]);
assert_eq!(m.check("YES, I agree"), MatchResult::Matched("YES".into()));
assert_eq!(m.check("Maybe"), MatchResult::NoMatch);
```
*/

/// Result of a prefix match check.
#[derive(Debug, Clone, PartialEq)]
pub enum MatchResult {
    /// The text started with this prefix.
    Matched(String),
    /// No configured prefix matched.
    NoMatch,
}

impl MatchResult {
    pub fn is_match(&self) -> bool { matches!(self, MatchResult::Matched(_)) }
    pub fn matched_prefix(&self) -> Option<&str> {
        if let MatchResult::Matched(s) = self { Some(s) } else { None }
    }
}

/// Match mode.
#[derive(Debug, Clone, PartialEq)]
pub enum MatchMode {
    CaseSensitive,
    CaseInsensitive,
}

/// Checks whether LLM outputs start with any of a set of expected prefixes.
#[derive(Debug, Clone, Default)]
pub struct PrefixMatcher {
    prefixes: Vec<String>,
    mode: Option<MatchMode>,
    trim_before_check: bool,
}

impl PrefixMatcher {
    pub fn new() -> Self { Self::default() }

    pub fn require(mut self, prefix: impl Into<String>) -> Self {
        self.prefixes.push(prefix.into()); self
    }

    pub fn require_any(mut self, prefixes: &[&str]) -> Self {
        self.prefixes.extend(prefixes.iter().map(|s| s.to_string())); self
    }

    pub fn case_insensitive(mut self) -> Self {
        self.mode = Some(MatchMode::CaseInsensitive); self
    }

    pub fn trim(mut self) -> Self { self.trim_before_check = true; self }

    /// Check if `text` starts with any configured prefix. Returns the first match.
    pub fn check(&self, text: &str) -> MatchResult {
        let candidate = if self.trim_before_check { text.trim_start() } else { text };
        let is_ci = self.mode == Some(MatchMode::CaseInsensitive);
        for prefix in &self.prefixes {
            let matches = if is_ci {
                candidate.to_lowercase().starts_with(&prefix.to_lowercase())
            } else {
                candidate.starts_with(prefix.as_str())
            };
            if matches {
                return MatchResult::Matched(prefix.clone());
            }
        }
        MatchResult::NoMatch
    }

    /// Strip the matched prefix from `text` and return the remainder, or None.
    pub fn strip(&self, text: &str) -> Option<String> {
        let candidate = if self.trim_before_check { text.trim_start() } else { text };
        let is_ci = self.mode == Some(MatchMode::CaseInsensitive);
        for prefix in &self.prefixes {
            let matches = if is_ci {
                candidate.to_lowercase().starts_with(&prefix.to_lowercase())
            } else {
                candidate.starts_with(prefix.as_str())
            };
            if matches {
                return Some(candidate[prefix.len()..].trim_start().to_string());
            }
        }
        None
    }

    /// True if text matches any prefix.
    pub fn is_valid(&self, text: &str) -> bool {
        self.check(text).is_match()
    }

    pub fn prefix_count(&self) -> usize { self.prefixes.len() }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn single_prefix_match() {
        let m = PrefixMatcher::new().require("YES");
        assert_eq!(m.check("YES I agree"), MatchResult::Matched("YES".into()));
    }

    #[test]
    fn single_prefix_no_match() {
        let m = PrefixMatcher::new().require("YES");
        assert_eq!(m.check("NO"), MatchResult::NoMatch);
    }

    #[test]
    fn multiple_prefixes_first_match_wins() {
        let m = PrefixMatcher::new().require("YES").require("NO");
        assert_eq!(m.check("YES: sure"), MatchResult::Matched("YES".into()));
    }

    #[test]
    fn require_any() {
        let m = PrefixMatcher::new().require_any(&["YES", "NO", "MAYBE"]);
        assert!(m.check("MAYBE later").is_match());
    }

    #[test]
    fn case_insensitive_match() {
        let m = PrefixMatcher::new().require("yes").case_insensitive();
        assert!(m.check("YES I agree").is_match());
    }

    #[test]
    fn case_sensitive_no_match() {
        let m = PrefixMatcher::new().require("yes");
        assert!(!m.check("YES").is_match());
    }

    #[test]
    fn trim_leading_whitespace() {
        let m = PrefixMatcher::new().require("OK").trim();
        assert!(m.check("   OK great").is_match());
    }

    #[test]
    fn strip_prefix_returns_remainder() {
        let m = PrefixMatcher::new().require("YES:");
        let rest = m.strip("YES: I agree");
        assert_eq!(rest.as_deref(), Some("I agree"));
    }

    #[test]
    fn strip_no_match_returns_none() {
        let m = PrefixMatcher::new().require("YES");
        assert!(m.strip("NO").is_none());
    }

    #[test]
    fn is_valid() {
        let m = PrefixMatcher::new().require("OK");
        assert!(m.is_valid("OK then"));
        assert!(!m.is_valid("bad"));
    }

    #[test]
    fn prefix_count() {
        let m = PrefixMatcher::new().require("A").require("B").require("C");
        assert_eq!(m.prefix_count(), 3);
    }

    #[test]
    fn empty_prefix_always_matches() {
        let m = PrefixMatcher::new().require("");
        assert!(m.check("anything").is_match());
    }

    #[test]
    fn matched_prefix_accessor() {
        let m = PrefixMatcher::new().require("YES");
        let res = m.check("YES!");
        assert_eq!(res.matched_prefix(), Some("YES"));
    }
}