Skip to main content

llama_cpp_bindings/
grammar_matcher.rs

1use std::panic::AssertUnwindSafe;
2use std::panic::catch_unwind;
3
4use llguidance::TokenParser;
5use llguidance::api::StopReason;
6
7use crate::error::grammar_runtime_error::GrammarRuntimeError;
8use crate::mask_outcome::MaskOutcome;
9
10enum StepOutcome<TValue> {
11    Produced(TValue),
12    BenignStop,
13}
14
15fn stop_reason_to_result(
16    stop_reason: StopReason,
17    detail: String,
18) -> Result<(), GrammarRuntimeError> {
19    match stop_reason {
20        StopReason::NotStopped
21        | StopReason::NoExtension
22        | StopReason::NoExtensionBias
23        | StopReason::EndOfSentence => Ok(()),
24        StopReason::InternalError => {
25            Err(GrammarRuntimeError::InternalParserError { message: detail })
26        }
27        StopReason::LexerTooComplex => {
28            Err(GrammarRuntimeError::LexerTooComplex { message: detail })
29        }
30        StopReason::ParserTooComplex => {
31            Err(GrammarRuntimeError::ParserTooComplex { message: detail })
32        }
33        StopReason::MaxTokensTotal | StopReason::MaxTokensParser => {
34            Err(GrammarRuntimeError::MaxTokensReached { message: detail })
35        }
36    }
37}
38
39pub struct GrammarMatcher {
40    parser: TokenParser,
41}
42
43impl GrammarMatcher {
44    #[must_use]
45    pub fn new(parser: TokenParser) -> Self {
46        let mut parser = parser;
47        if parser.is_fresh() {
48            parser.start_without_prompt();
49        }
50
51        Self { parser }
52    }
53
54    #[must_use]
55    pub fn deep_clone(&self) -> Self {
56        Self {
57            parser: self.parser.deep_clone(),
58        }
59    }
60
61    /// # Errors
62    /// Returns [`GrammarRuntimeError`] when the parser reaches a genuine error
63    /// state (distinct from a benign grammar completion).
64    pub fn compute_mask(&mut self) -> Result<MaskOutcome, GrammarRuntimeError> {
65        match self.run("compute_mask", TokenParser::compute_mask)? {
66            StepOutcome::Produced(mask) => Ok(MaskOutcome::Constrained(mask)),
67            StepOutcome::BenignStop => Ok(MaskOutcome::GrammarComplete),
68        }
69    }
70
71    /// # Errors
72    /// Returns [`GrammarRuntimeError`] when consuming the token drives the parser
73    /// into a genuine error state. A token that completes the grammar is a
74    /// benign stop and yields `Ok(())`.
75    pub fn consume_token(&mut self, token: u32) -> Result<(), GrammarRuntimeError> {
76        match self.run("consume_token", |parser| parser.consume_token(token))? {
77            StepOutcome::Produced(_) | StepOutcome::BenignStop => Ok(()),
78        }
79    }
80
81    /// # Errors
82    /// Returns [`GrammarRuntimeError`] when the parser cannot be reset.
83    pub fn reset(&mut self) -> Result<(), GrammarRuntimeError> {
84        match self.run("reset", TokenParser::reset)? {
85            StepOutcome::Produced(()) | StepOutcome::BenignStop => Ok(()),
86        }
87    }
88
89    fn run<TValue, TError>(
90        &mut self,
91        operation: &'static str,
92        op: impl FnOnce(&mut TokenParser) -> Result<TValue, TError>,
93    ) -> Result<StepOutcome<TValue>, GrammarRuntimeError> {
94        match catch_unwind(AssertUnwindSafe(|| op(&mut self.parser))) {
95            Ok(op_result) => {
96                if let Ok(value) = op_result {
97                    return Ok(StepOutcome::Produced(value));
98                }
99
100                let detail = self.parser.error_message().unwrap_or_default();
101                stop_reason_to_result(self.parser.stop_reason(), detail)?;
102
103                Ok(StepOutcome::BenignStop)
104            }
105            Err(_panic) => Err(GrammarRuntimeError::Panicked { operation }),
106        }
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use llguidance::api::StopReason;
113
114    use super::stop_reason_to_result;
115    use crate::error::grammar_runtime_error::GrammarRuntimeError;
116
117    #[test]
118    fn benign_stop_reasons_are_ok() {
119        for reason in [
120            StopReason::NotStopped,
121            StopReason::NoExtension,
122            StopReason::NoExtensionBias,
123            StopReason::EndOfSentence,
124        ] {
125            assert!(stop_reason_to_result(reason, String::new()).is_ok());
126        }
127    }
128
129    #[test]
130    fn internal_error_maps_to_internal_parser_error_with_message() {
131        assert_eq!(
132            stop_reason_to_result(StopReason::InternalError, "boom".to_string()),
133            Err(GrammarRuntimeError::InternalParserError {
134                message: "boom".to_string()
135            })
136        );
137    }
138
139    #[test]
140    fn lexer_too_complex_maps_to_lexer_too_complex() {
141        assert_eq!(
142            stop_reason_to_result(StopReason::LexerTooComplex, String::new()),
143            Err(GrammarRuntimeError::LexerTooComplex {
144                message: String::new()
145            })
146        );
147    }
148
149    #[test]
150    fn parser_too_complex_maps_to_parser_too_complex() {
151        assert_eq!(
152            stop_reason_to_result(StopReason::ParserTooComplex, String::new()),
153            Err(GrammarRuntimeError::ParserTooComplex {
154                message: String::new()
155            })
156        );
157    }
158
159    #[test]
160    fn max_token_stop_reasons_map_to_max_tokens_reached() {
161        for reason in [StopReason::MaxTokensTotal, StopReason::MaxTokensParser] {
162            assert_eq!(
163                stop_reason_to_result(reason, String::new()),
164                Err(GrammarRuntimeError::MaxTokensReached {
165                    message: String::new()
166                })
167            );
168        }
169    }
170}