/home/noah/src/realizar/src/grammar/mod.rs
Line | Count | Source |
1 | | //! Grammar-constrained generation for structured output |
2 | | //! |
3 | | //! Implements GBNF-style grammar constraints for LLM generation. |
4 | | //! Supports JSON schema validation and custom grammar rules. |
5 | | //! |
6 | | //! Reference: llama.cpp grammar implementation |
7 | | //! - GBNF format: Backus-Naur Form with extensions |
8 | | //! - Token masking: Efficiently filter invalid tokens |
9 | | //! - State machine: Track grammar state during generation |
10 | | |
11 | | use crate::error::{RealizarError, Result}; |
12 | | use serde::{Deserialize, Serialize}; |
13 | | use std::collections::{HashMap, HashSet}; |
14 | | |
15 | | // ============================================================================= |
16 | | // GRAMMAR RULE TYPES |
17 | | // ============================================================================= |
18 | | |
19 | | /// Grammar rule element types |
20 | | #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] |
21 | | pub enum GrammarElement { |
22 | | /// Literal character |
23 | | Char(char), |
24 | | /// Character range [a-z] |
25 | | CharRange(char, char), |
26 | | /// Reference to another rule |
27 | | RuleRef(String), |
28 | | /// Negated character set [^...] |
29 | | CharNot(Vec<char>), |
30 | | /// Any character |
31 | | Any, |
32 | | /// End of rule |
33 | | End, |
34 | | } |
35 | | |
36 | | /// Alternative production in a rule |
37 | | #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] |
38 | | pub struct GrammarAlternative { |
39 | | /// Sequence of elements in this alternative |
40 | | pub elements: Vec<GrammarElement>, |
41 | | } |
42 | | |
43 | | impl GrammarAlternative { |
44 | | /// Create new alternative from elements |
45 | 361 | pub fn new(elements: Vec<GrammarElement>) -> Self { |
46 | 361 | Self { elements } |
47 | 361 | } |
48 | | |
49 | | /// Create single-character alternative |
50 | 3 | pub fn char(c: char) -> Self { |
51 | 3 | Self { |
52 | 3 | elements: vec![GrammarElement::Char(c)], |
53 | 3 | } |
54 | 3 | } |
55 | | |
56 | | /// Check if this alternative is empty (epsilon production) |
57 | 2 | pub fn is_empty(&self) -> bool { |
58 | 2 | self.elements.is_empty() |
59 | 2 | } |
60 | | } |
61 | | |
62 | | /// Grammar rule with one or more alternatives |
63 | | #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] |
64 | | pub struct GrammarRule { |
65 | | /// Rule name |
66 | | pub name: String, |
67 | | /// Alternative productions |
68 | | pub alternatives: Vec<GrammarAlternative>, |
69 | | } |
70 | | |
71 | | impl GrammarRule { |
72 | | /// Create new rule with given alternatives |
73 | 151 | pub fn new(name: impl Into<String>, alternatives: Vec<GrammarAlternative>) -> Self { |
74 | 151 | Self { |
75 | 151 | name: name.into(), |
76 | 151 | alternatives, |
77 | 151 | } |
78 | 151 | } |
79 | | |
80 | | /// Create rule with single alternative |
81 | 42 | pub fn single(name: impl Into<String>, elements: Vec<GrammarElement>) -> Self { |
82 | 42 | Self { |
83 | 42 | name: name.into(), |
84 | 42 | alternatives: vec![GrammarAlternative::new(elements)], |
85 | 42 | } |
86 | 42 | } |
87 | | } |
88 | | |
89 | | // ============================================================================= |
90 | | // GRAMMAR DEFINITION |
91 | | // ============================================================================= |
92 | | |
93 | | /// Complete grammar definition |
94 | | #[derive(Debug, Clone, Default, Serialize, Deserialize)] |
95 | | pub struct Grammar { |
96 | | /// Rules by name |
97 | | rules: HashMap<String, GrammarRule>, |
98 | | /// Root rule name |
99 | | root: String, |
100 | | } |
101 | | |
102 | | impl Grammar { |
103 | | /// Create empty grammar |
104 | 4 | pub fn new() -> Self { |
105 | 4 | Self { |
106 | 4 | rules: HashMap::new(), |
107 | 4 | root: String::new(), |
108 | 4 | } |
109 | 4 | } |
110 | | |
111 | | /// Create grammar with root rule |
112 | 47 | pub fn with_root(root: impl Into<String>) -> Self { |
113 | 47 | Self { |
114 | 47 | rules: HashMap::new(), |
115 | 47 | root: root.into(), |
116 | 47 | } |
117 | 47 | } |
118 | | |
119 | | /// Add a rule to the grammar |
120 | 190 | pub fn add_rule(&mut self, rule: GrammarRule) { |
121 | 190 | if self.root.is_empty() { |
122 | 6 | self.root.clone_from(&rule.name); |
123 | 184 | } |
124 | 190 | self.rules.insert(rule.name.clone(), rule); |
125 | 190 | } |
126 | | |
127 | | /// Get rule by name |
128 | 345 | pub fn get_rule(&self, name: &str) -> Option<&GrammarRule> { |
129 | 345 | self.rules.get(name) |
130 | 345 | } |
131 | | |
132 | | /// Get root rule name |
133 | 127 | pub fn root(&self) -> &str { |
134 | 127 | &self.root |
135 | 127 | } |
136 | | |
137 | | /// Set root rule |
138 | 3 | pub fn set_root(&mut self, root: impl Into<String>) { |
139 | 3 | self.root = root.into(); |
140 | 3 | } |
141 | | |
142 | | /// Get all rule names |
143 | 2 | pub fn rule_names(&self) -> impl Iterator<Item = &String> { |
144 | 2 | self.rules.keys() |
145 | 2 | } |
146 | | |
147 | | /// Number of rules |
148 | 1 | pub fn len(&self) -> usize { |
149 | 1 | self.rules.len() |
150 | 1 | } |
151 | | |
152 | | /// Check if grammar is empty |
153 | 3 | pub fn is_empty(&self) -> bool { |
154 | 3 | self.rules.is_empty() |
155 | 3 | } |
156 | | |
157 | | /// Validate grammar has required rules |
158 | | /// |
159 | | /// # Errors |
160 | | /// |
161 | | /// Returns `InvalidConfiguration` if: |
162 | | /// - Grammar has no root rule |
163 | | /// - Root rule is not defined in the grammar |
164 | | /// - Any rule references an undefined rule |
165 | 42 | pub fn validate(&self) -> Result<()> { |
166 | 42 | if self.root.is_empty() { |
167 | 1 | return Err(RealizarError::InvalidConfiguration( |
168 | 1 | "Grammar has no root rule".to_string(), |
169 | 1 | )); |
170 | 41 | } |
171 | | |
172 | 41 | if !self.rules.contains_key(&self.root) { |
173 | 1 | return Err(RealizarError::InvalidConfiguration(format!( |
174 | 1 | "Root rule '{}' not found in grammar", |
175 | 1 | self.root |
176 | 1 | ))); |
177 | 40 | } |
178 | | |
179 | | // Check all rule references are valid |
180 | 75 | for rule in self.rules40 .values40 () { |
181 | 204 | for alt132 in &rule.alternatives { |
182 | 331 | for elem202 in &alt.elements { |
183 | 202 | if let GrammarElement::RuleRef(ref_name64 ) = elem { |
184 | 64 | if !self.rules.contains_key(ref_name) { |
185 | 3 | return Err(RealizarError::InvalidConfiguration(format!( |
186 | 3 | "Rule '{}' references undefined rule '{}'", |
187 | 3 | rule.name, ref_name |
188 | 3 | ))); |
189 | 61 | } |
190 | 138 | } |
191 | | } |
192 | | } |
193 | | } |
194 | | |
195 | 37 | Ok(()) |
196 | 42 | } |
197 | | } |
198 | | |
199 | | // ============================================================================= |
200 | | // GRAMMAR STATE MACHINE |
201 | | // ============================================================================= |
202 | | |
203 | | /// State in the grammar state machine |
204 | | #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] |
205 | | pub struct GrammarState { |
206 | | /// Current rule being matched |
207 | | pub rule: String, |
208 | | /// Alternative index within rule |
209 | | pub alt_idx: usize, |
210 | | /// Position within alternative |
211 | | pub elem_idx: usize, |
212 | | /// Stack of parent states (for rule references) |
213 | | pub stack: Vec<(String, usize, usize)>, |
214 | | } |
215 | | |
216 | | impl GrammarState { |
217 | | /// Create initial state from root rule |
218 | 42 | pub fn initial(root: &str) -> Self { |
219 | 42 | Self { |
220 | 42 | rule: root.to_string(), |
221 | 42 | alt_idx: 0, |
222 | 42 | elem_idx: 0, |
223 | 42 | stack: Vec::new(), |
224 | 42 | } |
225 | 42 | } |
226 | | |
227 | | /// Check if state is at end of rule |
228 | 18 | pub fn is_complete(&self, grammar: &Grammar) -> bool { |
229 | 18 | if let Some(rule17 ) = grammar.get_rule(&self.rule) { |
230 | 17 | if self.alt_idx < rule.alternatives.len() { |
231 | 16 | let alt = &rule.alternatives[self.alt_idx]; |
232 | 16 | return self.elem_idx >= alt.elements.len() && self.stack10 .is_empty10 (); |
233 | 1 | } |
234 | 1 | } |
235 | 2 | false |
236 | 18 | } |
237 | | |
238 | | /// Get current element being matched |
239 | 211 | pub fn current_element<'a>(&self, grammar: &'a Grammar) -> Option<&'a GrammarElement> { |
240 | 211 | grammar.get_rule(&self.rule).and_then(|rule| {210 |
241 | 210 | rule.alternatives |
242 | 210 | .get(self.alt_idx) |
243 | 210 | .and_then(|alt| alt.elements209 .get209 (self.elem_idx209 )) |
244 | 210 | }) |
245 | 211 | } |
246 | | } |
247 | | |
248 | | /// Grammar state machine for tracking generation progress |
249 | | #[derive(Debug, Clone)] |
250 | | pub struct GrammarStateMachine { |
251 | | /// The grammar being enforced |
252 | | grammar: Grammar, |
253 | | /// Current possible states (NFA-style) |
254 | | states: Vec<GrammarState>, |
255 | | /// Generated characters so far |
256 | | generated: String, |
257 | | } |
258 | | |
259 | | impl GrammarStateMachine { |
260 | | /// Create new state machine from grammar |
261 | | /// |
262 | | /// # Errors |
263 | | /// |
264 | | /// Returns an error if the grammar fails validation. |
265 | 36 | pub fn new(grammar: Grammar) -> Result<Self> { |
266 | 36 | grammar.validate()?0 ; |
267 | | |
268 | 36 | let initial = GrammarState::initial(grammar.root()); |
269 | 36 | let mut states = vec![initial]; |
270 | | |
271 | | // Expand initial states for all alternatives of root rule |
272 | 36 | if let Some(root_rule) = grammar.get_rule(grammar.root()) { |
273 | 36 | states.clear(); |
274 | 45 | for (alt_idx, _) in root_rule.alternatives.iter()36 .enumerate36 () { |
275 | 45 | states.push(GrammarState { |
276 | 45 | rule: grammar.root().to_string(), |
277 | 45 | alt_idx, |
278 | 45 | elem_idx: 0, |
279 | 45 | stack: Vec::new(), |
280 | 45 | }); |
281 | 45 | } |
282 | 0 | } |
283 | | |
284 | 36 | Ok(Self { |
285 | 36 | grammar, |
286 | 36 | states, |
287 | 36 | generated: String::new(), |
288 | 36 | }) |
289 | 36 | } |
290 | | |
291 | | /// Check if a character is valid at current state |
292 | 37 | pub fn is_valid_char(&self, c: char) -> bool { |
293 | 54 | for state43 in &self.states { |
294 | 43 | if self.can_accept_char(state, c) { |
295 | 26 | return true; |
296 | 17 | } |
297 | | } |
298 | 11 | false |
299 | 37 | } |
300 | | |
301 | | /// Get all valid characters at current state |
302 | 8 | pub fn valid_chars(&self) -> HashSet<char> { |
303 | 8 | let mut valid = HashSet::new(); |
304 | | |
305 | 18 | for state10 in &self.states { |
306 | 10 | self.collect_valid_chars(state, &mut valid); |
307 | 10 | } |
308 | | |
309 | 8 | valid |
310 | 8 | } |
311 | | |
312 | | /// Advance state machine with a character |
313 | 37 | pub fn advance(&mut self, c: char) -> bool { |
314 | 37 | let mut new_states = Vec::new(); |
315 | | |
316 | 77 | for state40 in &self.states { |
317 | 40 | if let Some(next_states31 ) = self.advance_state(state, c) { |
318 | 31 | new_states.extend(next_states); |
319 | 31 | }9 |
320 | | } |
321 | | |
322 | 37 | if new_states.is_empty() { |
323 | 7 | return false; |
324 | 30 | } |
325 | | |
326 | 30 | self.states = new_states; |
327 | 30 | self.generated.push(c); |
328 | 30 | true |
329 | 37 | } |
330 | | |
331 | | /// Check if generation is complete (valid end state) |
332 | 15 | pub fn is_complete(&self) -> bool { |
333 | 15 | self.states.iter().any(|s| s.is_complete(&self.grammar)) |
334 | 15 | } |
335 | | |
336 | | /// Check if any valid continuation exists |
337 | 2 | pub fn has_valid_continuation(&self) -> bool { |
338 | 2 | !self.states.is_empty() |
339 | 2 | } |
340 | | |
341 | | /// Get generated string so far |
342 | 3 | pub fn generated(&self) -> &str { |
343 | 3 | &self.generated |
344 | 3 | } |
345 | | |
346 | | /// Reset state machine |
347 | 2 | pub fn reset(&mut self) { |
348 | 2 | let initial = GrammarState::initial(self.grammar.root()); |
349 | 2 | self.states = vec![initial]; |
350 | | |
351 | | // Expand for all alternatives |
352 | 2 | if let Some(root_rule) = self.grammar.get_rule(self.grammar.root()) { |
353 | 2 | self.states.clear(); |
354 | 2 | for (alt_idx, _) in root_rule.alternatives.iter().enumerate() { |
355 | 2 | self.states.push(GrammarState { |
356 | 2 | rule: self.grammar.root().to_string(), |
357 | 2 | alt_idx, |
358 | 2 | elem_idx: 0, |
359 | 2 | stack: Vec::new(), |
360 | 2 | }); |
361 | 2 | } |
362 | 0 | } |
363 | | |
364 | 2 | self.generated.clear(); |
365 | 2 | } |
366 | | |
367 | | // Internal: Check if state can accept character |
368 | 154 | fn can_accept_char(&self, state: &GrammarState, c: char) -> bool { |
369 | 154 | if let Some(elem) = state.current_element(&self.grammar) { |
370 | 154 | match elem { |
371 | 22 | GrammarElement::Char(expected) => c == *expected, |
372 | 12 | GrammarElement::CharRange(start, end) => c >= *start && c <= *end6 , |
373 | 100 | GrammarElement::CharNot(excluded) => !excluded.contains(&c), |
374 | 4 | GrammarElement::Any => true, |
375 | 14 | GrammarElement::RuleRef(rule_name) => { |
376 | | // Need to check if any alternative of referenced rule accepts c |
377 | 14 | if let Some(rule) = self.grammar.get_rule(rule_name) { |
378 | 16 | for (alt_idx, _) in rule.alternatives.iter()14 .enumerate14 () { |
379 | 16 | let sub_state = GrammarState { |
380 | 16 | rule: rule_name.clone(), |
381 | 16 | alt_idx, |
382 | 16 | elem_idx: 0, |
383 | 16 | stack: Vec::new(), |
384 | 16 | }; |
385 | 16 | if self.can_accept_char(&sub_state, c) { |
386 | 8 | return true; |
387 | 8 | } |
388 | | } |
389 | 0 | } |
390 | 6 | false |
391 | | }, |
392 | 2 | GrammarElement::End => false, |
393 | | } |
394 | | } else { |
395 | 0 | false |
396 | | } |
397 | 154 | } |
398 | | |
399 | | // Internal: Collect valid characters for a state |
400 | 10 | fn collect_valid_chars(&self, state: &GrammarState, valid: &mut HashSet<char>) { |
401 | 10 | if let Some(elem) = state.current_element(&self.grammar) { |
402 | 10 | match elem { |
403 | 6 | GrammarElement::Char(c) => { |
404 | 6 | valid.insert(*c); |
405 | 6 | }, |
406 | 1 | GrammarElement::CharRange(start, end) => { |
407 | 3 | for c in *start1 ..=*end1 { |
408 | 3 | valid.insert(c); |
409 | 3 | } |
410 | | }, |
411 | 1 | GrammarElement::CharNot(_excluded) => { |
412 | | // For negated sets, we'd need to add all chars except excluded |
413 | | // This is expensive, so for now we mark as "any printable" |
414 | 96 | for c95 in ' '..='~' { |
415 | 95 | if self.can_accept_char(state, c) { |
416 | 92 | valid.insert(c); |
417 | 92 | }3 |
418 | | } |
419 | | }, |
420 | | GrammarElement::Any => { |
421 | | // Add common printable characters |
422 | 96 | for c95 in ' '..='~' { |
423 | 95 | valid.insert(c); |
424 | 95 | } |
425 | | }, |
426 | 0 | GrammarElement::RuleRef(rule_name) => { |
427 | | // Recurse into referenced rule |
428 | 0 | if let Some(rule) = self.grammar.get_rule(rule_name) { |
429 | 0 | for (alt_idx, _) in rule.alternatives.iter().enumerate() { |
430 | 0 | let sub_state = GrammarState { |
431 | 0 | rule: rule_name.clone(), |
432 | 0 | alt_idx, |
433 | 0 | elem_idx: 0, |
434 | 0 | stack: Vec::new(), |
435 | 0 | }; |
436 | 0 | self.collect_valid_chars(&sub_state, valid); |
437 | 0 | } |
438 | 0 | } |
439 | | }, |
440 | 1 | GrammarElement::End => {}, |
441 | | } |
442 | 0 | } |
443 | 10 | } |
444 | | |
445 | | // Internal: Advance state and return new states |
446 | 44 | fn advance_state(&self, state: &GrammarState, c: char) -> Option<Vec<GrammarState>> { |
447 | 44 | let elem = state.current_element(&self.grammar)?0 ; |
448 | | |
449 | 44 | match elem { |
450 | 33 | GrammarElement::Char(expected) => { |
451 | 33 | if c == *expected { |
452 | 27 | Some(vec![self.next_state(state)]) |
453 | | } else { |
454 | 6 | None |
455 | | } |
456 | | }, |
457 | 3 | GrammarElement::CharRange(start, end) => { |
458 | 3 | if c >= *start && c <= *end { |
459 | 2 | Some(vec![self.next_state(state)]) |
460 | | } else { |
461 | 1 | None |
462 | | } |
463 | | }, |
464 | 2 | GrammarElement::CharNot(excluded) => { |
465 | 2 | if !excluded.contains(&c) { |
466 | 1 | Some(vec![self.next_state(state)]) |
467 | | } else { |
468 | 1 | None |
469 | | } |
470 | | }, |
471 | 1 | GrammarElement::Any => Some(vec![self.next_state(state)]), |
472 | 4 | GrammarElement::RuleRef(rule_name) => { |
473 | | // Enter referenced rule |
474 | 4 | let rule = self.grammar.get_rule(rule_name)?0 ; |
475 | 4 | let mut new_states = Vec::new(); |
476 | | |
477 | 4 | for (alt_idx, _) in rule.alternatives.iter().enumerate() { |
478 | 4 | let mut sub_state = GrammarState { |
479 | 4 | rule: rule_name.clone(), |
480 | 4 | alt_idx, |
481 | 4 | elem_idx: 0, |
482 | 4 | stack: state.stack.clone(), |
483 | 4 | }; |
484 | | // Push return address |
485 | 4 | sub_state |
486 | 4 | .stack |
487 | 4 | .push((state.rule.clone(), state.alt_idx, state.elem_idx + 1)); |
488 | | |
489 | 4 | if let Some(advanced3 ) = self.advance_state(&sub_state, c) { |
490 | 3 | new_states.extend(advanced); |
491 | 3 | }1 |
492 | | } |
493 | | |
494 | 4 | if new_states.is_empty() { |
495 | 1 | None |
496 | | } else { |
497 | 3 | Some(new_states) |
498 | | } |
499 | | }, |
500 | 1 | GrammarElement::End => None, |
501 | | } |
502 | 44 | } |
503 | | |
504 | | // Internal: Create next state after consuming element |
505 | 31 | fn next_state(&self, state: &GrammarState) -> GrammarState { |
506 | 31 | let mut new_state = state.clone(); |
507 | 31 | new_state.elem_idx += 1; |
508 | | |
509 | | // Check if we've completed current alternative |
510 | 31 | if let Some(rule) = self.grammar.get_rule(&state.rule) { |
511 | 31 | if let Some(alt) = rule.alternatives.get(state.alt_idx) { |
512 | 31 | if new_state.elem_idx >= alt.elements.len() { |
513 | | // Pop from stack if there's a return address |
514 | 15 | if let Some((ret_rule2 , ret_alt2 , ret_elem2 )) = new_state.stack.pop() { |
515 | 2 | new_state.rule = ret_rule; |
516 | 2 | new_state.alt_idx = ret_alt; |
517 | 2 | new_state.elem_idx = ret_elem; |
518 | 13 | } |
519 | 16 | } |
520 | 0 | } |
521 | 0 | } |
522 | | |
523 | 31 | new_state |
524 | 31 | } |
525 | | } |
526 | | |
527 | | // ============================================================================= |
528 | | // JSON SCHEMA GRAMMAR BUILDER |
529 | | // ============================================================================= |
530 | | |
531 | | /// JSON Schema types for grammar generation |
532 | | #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] |
533 | | pub enum JsonSchemaType { |
534 | | /// String type |
535 | | String, |
536 | | /// Integer type |
537 | | Integer, |
538 | | /// Number type (float) |
539 | | Number, |
540 | | /// Boolean type |
541 | | Boolean, |
542 | | /// Null type |
543 | | Null, |
544 | | /// Array type with item schema |
545 | | Array(Box<JsonSchemaType>), |
546 | | /// Object type with properties |
547 | | Object(Vec<(String, JsonSchemaType, bool)>), // (name, type, required) |
548 | | /// Enum with allowed values |
549 | | Enum(Vec<String>), |
550 | | /// Any type |
551 | | Any, |
552 | | } |
553 | | |
554 | | /// Build grammar from JSON schema type |
555 | 12 | pub fn grammar_from_json_schema(schema: &JsonSchemaType) -> Grammar { |
556 | 12 | let mut grammar = Grammar::with_root("root"); |
557 | | |
558 | | // Add whitespace rule |
559 | 12 | grammar.add_rule(GrammarRule::new( |
560 | | "ws", |
561 | 12 | vec![ |
562 | 12 | GrammarAlternative::new(vec![]), // Empty (epsilon) |
563 | 12 | GrammarAlternative::new(vec![ |
564 | 12 | GrammarElement::Char(' '), |
565 | 12 | GrammarElement::RuleRef("ws".to_string()), |
566 | | ]), |
567 | 12 | GrammarAlternative::new(vec![ |
568 | 12 | GrammarElement::Char('\n'), |
569 | 12 | GrammarElement::RuleRef("ws".to_string()), |
570 | | ]), |
571 | 12 | GrammarAlternative::new(vec![ |
572 | 12 | GrammarElement::Char('\t'), |
573 | 12 | GrammarElement::RuleRef("ws".to_string()), |
574 | | ]), |
575 | | ], |
576 | | )); |
577 | | |
578 | | // Add digit rule |
579 | 12 | grammar.add_rule(GrammarRule::new( |
580 | | "digit", |
581 | 12 | vec![GrammarAlternative::new(vec![GrammarElement::CharRange( |
582 | 12 | '0', '9', |
583 | 12 | )])], |
584 | | )); |
585 | | |
586 | | // Add digits rule (one or more) |
587 | 12 | grammar.add_rule(GrammarRule::new( |
588 | | "digits", |
589 | 12 | vec![ |
590 | 12 | GrammarAlternative::new(vec![GrammarElement::RuleRef("digit".to_string())]), |
591 | 12 | GrammarAlternative::new(vec![ |
592 | 12 | GrammarElement::RuleRef("digit".to_string()), |
593 | 12 | GrammarElement::RuleRef("digits".to_string()), |
594 | | ]), |
595 | | ], |
596 | | )); |
597 | | |
598 | | // Add string character rule |
599 | 12 | grammar.add_rule(GrammarRule::new( |
600 | | "string_char", |
601 | 12 | vec![ |
602 | 12 | GrammarAlternative::new(vec![GrammarElement::CharNot(vec!['"', '\\', '\n'])]), |
603 | 12 | GrammarAlternative::new(vec![GrammarElement::Char('\\'), GrammarElement::Char('"')]), |
604 | 12 | GrammarAlternative::new(vec![GrammarElement::Char('\\'), GrammarElement::Char('\\')]), |
605 | 12 | GrammarAlternative::new(vec![GrammarElement::Char('\\'), GrammarElement::Char('n')]), |
606 | | ], |
607 | | )); |
608 | | |
609 | | // Add string content rule |
610 | 12 | grammar.add_rule(GrammarRule::new( |
611 | | "string_content", |
612 | 12 | vec![ |
613 | 12 | GrammarAlternative::new(vec![]), // Empty |
614 | 12 | GrammarAlternative::new(vec![ |
615 | 12 | GrammarElement::RuleRef("string_char".to_string()), |
616 | 12 | GrammarElement::RuleRef("string_content".to_string()), |
617 | | ]), |
618 | | ], |
619 | | )); |
620 | | |
621 | | // Add base type rules based on schema |
622 | 12 | add_schema_rules(&mut grammar, "root", schema); |
623 | | |
624 | 12 | grammar |
625 | 12 | } |
626 | | |
627 | 24 | fn add_schema_rules(grammar: &mut Grammar, rule_name: &str, schema: &JsonSchemaType) { |
628 | 24 | match schema { |
629 | 4 | JsonSchemaType::String => { |
630 | 4 | grammar.add_rule(GrammarRule::new( |
631 | 4 | rule_name, |
632 | 4 | vec![GrammarAlternative::new(vec![ |
633 | 4 | GrammarElement::Char('"'), |
634 | 4 | GrammarElement::RuleRef("string_content".to_string()), |
635 | 4 | GrammarElement::Char('"'), |
636 | 4 | ])], |
637 | 4 | )); |
638 | 4 | }, |
639 | 5 | JsonSchemaType::Integer => { |
640 | 5 | grammar.add_rule(GrammarRule::new( |
641 | 5 | rule_name, |
642 | 5 | vec![ |
643 | 5 | GrammarAlternative::new(vec![GrammarElement::RuleRef("digits".to_string())]), |
644 | 5 | GrammarAlternative::new(vec![ |
645 | 5 | GrammarElement::Char('-'), |
646 | 5 | GrammarElement::RuleRef("digits".to_string()), |
647 | 5 | ]), |
648 | 5 | ], |
649 | 5 | )); |
650 | 5 | }, |
651 | 2 | JsonSchemaType::Number => { |
652 | 2 | // integer or decimal |
653 | 2 | grammar.add_rule(GrammarRule::new( |
654 | 2 | rule_name, |
655 | 2 | vec![ |
656 | 2 | GrammarAlternative::new(vec![GrammarElement::RuleRef("digits".to_string())]), |
657 | 2 | GrammarAlternative::new(vec![ |
658 | 2 | GrammarElement::Char('-'), |
659 | 2 | GrammarElement::RuleRef("digits".to_string()), |
660 | 2 | ]), |
661 | 2 | GrammarAlternative::new(vec![ |
662 | 2 | GrammarElement::RuleRef("digits".to_string()), |
663 | 2 | GrammarElement::Char('.'), |
664 | 2 | GrammarElement::RuleRef("digits".to_string()), |
665 | 2 | ]), |
666 | 2 | GrammarAlternative::new(vec![ |
667 | 2 | GrammarElement::Char('-'), |
668 | 2 | GrammarElement::RuleRef("digits".to_string()), |
669 | 2 | GrammarElement::Char('.'), |
670 | 2 | GrammarElement::RuleRef("digits".to_string()), |
671 | 2 | ]), |
672 | 2 | ], |
673 | 2 | )); |
674 | 2 | }, |
675 | 3 | JsonSchemaType::Boolean => { |
676 | 3 | grammar.add_rule(GrammarRule::new( |
677 | 3 | rule_name, |
678 | 3 | vec![ |
679 | 3 | GrammarAlternative::new(vec![ |
680 | 3 | GrammarElement::Char('t'), |
681 | 3 | GrammarElement::Char('r'), |
682 | 3 | GrammarElement::Char('u'), |
683 | 3 | GrammarElement::Char('e'), |
684 | 3 | ]), |
685 | 3 | GrammarAlternative::new(vec![ |
686 | 3 | GrammarElement::Char('f'), |
687 | 3 | GrammarElement::Char('a'), |
688 | 3 | GrammarElement::Char('l'), |
689 | 3 | GrammarElement::Char('s'), |
690 | 3 | GrammarElement::Char('e'), |
691 | 3 | ]), |
692 | 3 | ], |
693 | 3 | )); |
694 | 3 | }, |
695 | 2 | JsonSchemaType::Null => { |
696 | 2 | grammar.add_rule(GrammarRule::new( |
697 | 2 | rule_name, |
698 | 2 | vec![GrammarAlternative::new(vec![ |
699 | 2 | GrammarElement::Char('n'), |
700 | 2 | GrammarElement::Char('u'), |
701 | 2 | GrammarElement::Char('l'), |
702 | 2 | GrammarElement::Char('l'), |
703 | 2 | ])], |
704 | 2 | )); |
705 | 2 | }, |
706 | 1 | JsonSchemaType::Enum(values) => { |
707 | 1 | let alternatives: Vec<GrammarAlternative> = values |
708 | 1 | .iter() |
709 | 2 | .map1 (|v| { |
710 | 2 | let mut elements = vec![GrammarElement::Char('"')]; |
711 | 7 | for c in v2 .chars2 () { |
712 | 7 | elements.push(GrammarElement::Char(c)); |
713 | 7 | } |
714 | 2 | elements.push(GrammarElement::Char('"')); |
715 | 2 | GrammarAlternative::new(elements) |
716 | 2 | }) |
717 | 1 | .collect(); |
718 | 1 | grammar.add_rule(GrammarRule::new(rule_name, alternatives)); |
719 | | }, |
720 | 3 | JsonSchemaType::Array(item_schema) => { |
721 | 3 | let item_rule = format!("{rule_name}_item"); |
722 | 3 | add_schema_rules(grammar, &item_rule, item_schema); |
723 | 3 | |
724 | 3 | let items_rule = format!("{rule_name}_items"); |
725 | 3 | grammar.add_rule(GrammarRule::new( |
726 | 3 | &items_rule, |
727 | 3 | vec![ |
728 | 3 | GrammarAlternative::new(vec![]), // Empty |
729 | 3 | GrammarAlternative::new(vec![ |
730 | 3 | GrammarElement::Char(','), |
731 | 3 | GrammarElement::RuleRef("ws".to_string()), |
732 | 3 | GrammarElement::RuleRef(item_rule.clone()), |
733 | 3 | GrammarElement::RuleRef(items_rule.clone()), |
734 | 3 | ]), |
735 | 3 | ], |
736 | 3 | )); |
737 | 3 | |
738 | 3 | grammar.add_rule(GrammarRule::new( |
739 | 3 | rule_name, |
740 | 3 | vec![ |
741 | 3 | // Empty array |
742 | 3 | GrammarAlternative::new(vec![ |
743 | 3 | GrammarElement::Char('['), |
744 | 3 | GrammarElement::RuleRef("ws".to_string()), |
745 | 3 | GrammarElement::Char(']'), |
746 | 3 | ]), |
747 | 3 | // Non-empty array |
748 | 3 | GrammarAlternative::new(vec![ |
749 | 3 | GrammarElement::Char('['), |
750 | 3 | GrammarElement::RuleRef("ws".to_string()), |
751 | 3 | GrammarElement::RuleRef(item_rule), |
752 | 3 | GrammarElement::RuleRef(items_rule), |
753 | 3 | GrammarElement::RuleRef("ws".to_string()), |
754 | 3 | GrammarElement::Char(']'), |
755 | 3 | ]), |
756 | 3 | ], |
757 | 3 | )); |
758 | 3 | }, |
759 | 3 | JsonSchemaType::Object(properties) => { |
760 | | // Build object grammar with properties |
761 | 3 | if properties.is_empty() { |
762 | 1 | // Empty object |
763 | 1 | grammar.add_rule(GrammarRule::new( |
764 | 1 | rule_name, |
765 | 1 | vec![GrammarAlternative::new(vec![ |
766 | 1 | GrammarElement::Char('{'), |
767 | 1 | GrammarElement::RuleRef("ws".to_string()), |
768 | 1 | GrammarElement::Char('}'), |
769 | 1 | ])], |
770 | 1 | )); |
771 | 1 | } else { |
772 | | // Object with properties |
773 | 2 | let mut elements = vec![ |
774 | 2 | GrammarElement::Char('{'), |
775 | 2 | GrammarElement::RuleRef("ws".to_string()), |
776 | | ]; |
777 | | |
778 | 5 | for (i, (prop_name, prop_type, _required)) in properties.iter()2 .enumerate2 () { |
779 | 5 | if i > 0 { |
780 | 3 | elements.push(GrammarElement::Char(',')); |
781 | 3 | elements.push(GrammarElement::RuleRef("ws".to_string())); |
782 | 3 | }2 |
783 | | |
784 | | // Property name |
785 | 5 | elements.push(GrammarElement::Char('"')); |
786 | 22 | for c in prop_name5 .chars5 () { |
787 | 22 | elements.push(GrammarElement::Char(c)); |
788 | 22 | } |
789 | 5 | elements.push(GrammarElement::Char('"')); |
790 | 5 | elements.push(GrammarElement::RuleRef("ws".to_string())); |
791 | 5 | elements.push(GrammarElement::Char(':')); |
792 | 5 | elements.push(GrammarElement::RuleRef("ws".to_string())); |
793 | | |
794 | | // Property value |
795 | 5 | let prop_rule = format!("{rule_name}_{prop_name}"); |
796 | 5 | add_schema_rules(grammar, &prop_rule, prop_type); |
797 | 5 | elements.push(GrammarElement::RuleRef(prop_rule)); |
798 | | } |
799 | | |
800 | 2 | elements.push(GrammarElement::RuleRef("ws".to_string())); |
801 | 2 | elements.push(GrammarElement::Char('}')); |
802 | | |
803 | 2 | grammar.add_rule(GrammarRule::new( |
804 | 2 | rule_name, |
805 | 2 | vec![GrammarAlternative::new(elements)], |
806 | | )); |
807 | | } |
808 | | }, |
809 | | JsonSchemaType::Any => { |
810 | | // Any JSON value |
811 | 1 | grammar.add_rule(GrammarRule::new( |
812 | 1 | rule_name, |
813 | 1 | vec![ |
814 | 1 | GrammarAlternative::new(vec![GrammarElement::RuleRef( |
815 | 1 | "string_value".to_string(), |
816 | 1 | )]), |
817 | 1 | GrammarAlternative::new(vec![GrammarElement::RuleRef("number".to_string())]), |
818 | 1 | GrammarAlternative::new(vec![GrammarElement::RuleRef("boolean".to_string())]), |
819 | 1 | GrammarAlternative::new(vec![GrammarElement::RuleRef("null".to_string())]), |
820 | | ], |
821 | | )); |
822 | | |
823 | | // Add helper rules if not present |
824 | 1 | if grammar.get_rule("string_value").is_none() { |
825 | 1 | add_schema_rules(grammar, "string_value", &JsonSchemaType::String); |
826 | 1 | }0 |
827 | 1 | if grammar.get_rule("number").is_none() { |
828 | 1 | add_schema_rules(grammar, "number", &JsonSchemaType::Number); |
829 | 1 | }0 |
830 | 1 | if grammar.get_rule("boolean").is_none() { |
831 | 1 | add_schema_rules(grammar, "boolean", &JsonSchemaType::Boolean); |
832 | 1 | }0 |
833 | 1 | if grammar.get_rule("null").is_none() { |
834 | 1 | add_schema_rules(grammar, "null", &JsonSchemaType::Null); |
835 | 1 | }0 |
836 | | }, |
837 | | } |
838 | 24 | } |
839 | | |
840 | | // ============================================================================= |
841 | | // TOKEN MASKING FOR CONSTRAINED GENERATION |
842 | | // ============================================================================= |
843 | | |
844 | | /// Token mask for constrained generation |
845 | | #[derive(Debug, Clone)] |
846 | | pub struct TokenMask { |
847 | | /// Allowed token IDs |
848 | | pub allowed: HashSet<u32>, |
849 | | /// Whether to allow end-of-sequence |
850 | | pub allow_eos: bool, |
851 | | } |
852 | | |
853 | | impl TokenMask { |
854 | | /// Create mask allowing all tokens |
855 | 1 | pub fn allow_all(vocab_size: usize) -> Self { |
856 | 1 | Self { |
857 | 1 | allowed: (0..vocab_size as u32).collect(), |
858 | 1 | allow_eos: true, |
859 | 1 | } |
860 | 1 | } |
861 | | |
862 | | /// Create mask from allowed set |
863 | 7 | pub fn from_allowed(allowed: HashSet<u32>, allow_eos: bool) -> Self { |
864 | 7 | Self { allowed, allow_eos } |
865 | 7 | } |
866 | | |
867 | | /// Check if token is allowed |
868 | 17 | pub fn is_allowed(&self, token_id: u32) -> bool { |
869 | 17 | self.allowed.contains(&token_id) |
870 | 17 | } |
871 | | |
872 | | /// Apply mask to logits (set disallowed to -inf) |
873 | 1 | pub fn apply_to_logits(&self, logits: &mut [f32]) { |
874 | 5 | for (i, logit) in logits1 .iter_mut1 ().enumerate1 () { |
875 | 5 | if !self.allowed.contains(&(i as u32)) { |
876 | 3 | *logit = f32::NEG_INFINITY; |
877 | 3 | }2 |
878 | | } |
879 | 1 | } |
880 | | |
881 | | /// Number of allowed tokens |
882 | 2 | pub fn num_allowed(&self) -> usize { |
883 | 2 | self.allowed.len() |
884 | 2 | } |
885 | | } |
886 | | |
887 | | /// Grammar-based token masker |
888 | | pub struct GrammarTokenMasker { |
889 | | /// State machine tracking grammar state |
890 | | state_machine: GrammarStateMachine, |
891 | | /// Token to string mapping |
892 | | token_strings: HashMap<u32, String>, |
893 | | /// EOS token ID |
894 | | eos_token_id: u32, |
895 | | } |
896 | | |
897 | | impl GrammarTokenMasker { |
898 | | /// Create new masker from grammar and vocabulary |
899 | | /// |
900 | | /// # Errors |
901 | | /// |
902 | | /// Returns an error if the grammar fails validation. |
903 | 7 | pub fn new( |
904 | 7 | grammar: Grammar, |
905 | 7 | token_strings: HashMap<u32, String>, |
906 | 7 | eos_token_id: u32, |
907 | 7 | ) -> Result<Self> { |
908 | 7 | let state_machine = GrammarStateMachine::new(grammar)?0 ; |
909 | 7 | Ok(Self { |
910 | 7 | state_machine, |
911 | 7 | token_strings, |
912 | 7 | eos_token_id, |
913 | 7 | }) |
914 | 7 | } |
915 | | |
916 | | /// Get mask for current state |
917 | 4 | pub fn get_mask(&self) -> TokenMask { |
918 | 4 | let valid_chars = self.state_machine.valid_chars(); |
919 | 4 | let mut allowed = HashSet::new(); |
920 | | |
921 | 14 | for (token_id10 , token_str10 ) in &self.token_strings { |
922 | | // Check if token's first character is valid |
923 | 10 | if let Some(first_char9 ) = token_str.chars().next() { |
924 | 9 | if valid_chars.contains(&first_char) { |
925 | | // For single-char tokens, this is sufficient |
926 | | // For multi-char tokens, we'd need to simulate all chars |
927 | 3 | if token_str.len() == 1 { |
928 | 2 | allowed.insert(*token_id); |
929 | 2 | } else { |
930 | | // Check if all characters in token are valid sequence |
931 | 1 | let mut temp_sm = self.state_machine.clone(); |
932 | 1 | let mut all_valid = true; |
933 | 2 | for c in token_str1 .chars1 () { |
934 | 2 | if !temp_sm.advance(c) { |
935 | 0 | all_valid = false; |
936 | 0 | break; |
937 | 2 | } |
938 | | } |
939 | 1 | if all_valid { |
940 | 1 | allowed.insert(*token_id); |
941 | 1 | }0 |
942 | | } |
943 | 6 | } |
944 | 1 | } |
945 | | } |
946 | | |
947 | 4 | let allow_eos = self.state_machine.is_complete(); |
948 | | |
949 | 4 | TokenMask::from_allowed(allowed, allow_eos) |
950 | 4 | } |
951 | | |
952 | | /// Advance masker with selected token |
953 | 5 | pub fn advance_token(&mut self, token_id: u32) -> bool { |
954 | 5 | if let Some(token_str4 ) = self.token_strings.get(&token_id) { |
955 | 4 | for c in token_str.chars() { |
956 | 4 | if !self.state_machine.advance(c) { |
957 | 1 | return false; |
958 | 3 | } |
959 | | } |
960 | 3 | true |
961 | | } else { |
962 | 1 | false |
963 | | } |
964 | 5 | } |
965 | | |
966 | | /// Check if generation is complete |
967 | 3 | pub fn is_complete(&self) -> bool { |
968 | 3 | self.state_machine.is_complete() |
969 | 3 | } |
970 | | |
971 | | /// Reset masker state |
972 | 1 | pub fn reset(&mut self) { |
973 | 1 | self.state_machine.reset(); |
974 | 1 | } |
975 | | |
976 | | /// Get EOS token ID |
977 | 1 | pub fn eos_token_id(&self) -> u32 { |
978 | 1 | self.eos_token_id |
979 | 1 | } |
980 | | } |
981 | | |
982 | | // ============================================================================= |
983 | | // TOOL CALLING / FUNCTION CALLING |
984 | | // ============================================================================= |
985 | | // |
986 | | // Implements OpenAI-style tool/function calling for LLM inference. |
987 | | // Allows models to generate structured function calls that can be executed |
988 | | // and results fed back into the conversation. |
989 | | // |
990 | | // Reference: OpenAI Function Calling API |
991 | | // - Tool definitions with JSON Schema parameters |
992 | | // - Tool choice: auto, required, none, or specific tool |
993 | | // - Tool call parsing from model output |
994 | | // - Grammar generation for constrained tool output |
995 | | |
996 | | /// JSON Schema property type for tool parameters |
997 | | #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] |
998 | | #[serde(rename_all = "snake_case")] |
999 | | pub enum ToolParameterType { |
1000 | | /// String parameter |
1001 | | #[default] |
1002 | | String, |
1003 | | /// Integer parameter |
1004 | | Integer, |
1005 | | /// Number parameter (float) |
1006 | | Number, |
1007 | | /// Boolean parameter |
1008 | | Boolean, |
1009 | | /// Array parameter |
1010 | | Array { |
1011 | | /// Type of array items |
1012 | | items: Box<ToolParameterType>, |
1013 | | }, |
1014 | | /// Object parameter with properties |
1015 | | Object { |
1016 | | /// Properties of the object |
1017 | | properties: Vec<ToolParameter>, |
1018 | | }, |
1019 | | /// Enum of allowed string values |
1020 | | Enum(Vec<String>), |
1021 | | } |
1022 | | |
1023 | | /// Tool parameter definition |
1024 | | #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] |
1025 | | pub struct ToolParameter { |
1026 | | /// Parameter name |
1027 | | pub name: String, |
1028 | | /// Parameter description |
1029 | | pub description: String, |
1030 | | /// Parameter type |
1031 | | #[serde(rename = "type")] |
1032 | | pub param_type: ToolParameterType, |
1033 | | /// Whether the parameter is required |
1034 | | #[serde(default)] |
1035 | | pub required: bool, |
1036 | | /// Default value (JSON string) |
1037 | | #[serde(skip_serializing_if = "Option::is_none")] |
1038 | | pub default: Option<String>, |
1039 | | } |
1040 | | |
1041 | | impl ToolParameter { |
1042 | | /// Create new required string parameter |
1043 | 10 | pub fn required_string(name: impl Into<String>, description: impl Into<String>) -> Self { |
1044 | 10 | Self { |
1045 | 10 | name: name.into(), |
1046 | 10 | description: description.into(), |
1047 | 10 | param_type: ToolParameterType::String, |
1048 | 10 | required: true, |
1049 | 10 | default: None, |
1050 | 10 | } |
1051 | 10 | } |
1052 | | |
1053 | | /// Create new optional string parameter |
1054 | 6 | pub fn optional_string(name: impl Into<String>, description: impl Into<String>) -> Self { |
1055 | 6 | Self { |
1056 | 6 | name: name.into(), |
1057 | 6 | description: description.into(), |
1058 | 6 | param_type: ToolParameterType::String, |
1059 | 6 | required: false, |
1060 | 6 | default: None, |
1061 | 6 | } |
1062 | 6 | } |
1063 | | |
1064 | | /// Create new required integer parameter |
1065 | 2 | pub fn required_int(name: impl Into<String>, description: impl Into<String>) -> Self { |
1066 | 2 | Self { |
1067 | 2 | name: name.into(), |
1068 | 2 | description: description.into(), |
1069 | 2 | param_type: ToolParameterType::Integer, |
1070 | 2 | required: true, |
1071 | 2 | default: None, |
1072 | 2 | } |
1073 | 2 | } |
1074 | | |
1075 | | /// Create new enum parameter |
1076 | 1 | pub fn required_enum( |
1077 | 1 | name: impl Into<String>, |
1078 | 1 | description: impl Into<String>, |
1079 | 1 | values: Vec<String>, |
1080 | 1 | ) -> Self { |
1081 | 1 | Self { |
1082 | 1 | name: name.into(), |
1083 | 1 | description: description.into(), |
1084 | 1 | param_type: ToolParameterType::Enum(values), |
1085 | 1 | required: true, |
1086 | 1 | default: None, |
1087 | 1 | } |
1088 | 1 | } |
1089 | | |
1090 | | /// Set default value |
1091 | | #[must_use] |
1092 | 1 | pub fn with_default(mut self, default: impl Into<String>) -> Self { |
1093 | 1 | self.default = Some(default.into()); |
1094 | 1 | self |
1095 | 1 | } |
1096 | | } |
1097 | | |
1098 | | /// Tool/function definition |
1099 | | #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] |
1100 | | pub struct ToolDefinition { |
1101 | | /// Tool name (must be valid identifier: [a-zA-Z_][a-zA-Z0-9_]*) |
1102 | | pub name: String, |
1103 | | /// Tool description |
1104 | | pub description: String, |
1105 | | /// Tool parameters |
1106 | | pub parameters: Vec<ToolParameter>, |
1107 | | } |
1108 | | |
1109 | | impl ToolDefinition { |
1110 | | /// Create new tool definition |
1111 | 35 | pub fn new( |
1112 | 35 | name: impl Into<String>, |
1113 | 35 | description: impl Into<String>, |
1114 | 35 | parameters: Vec<ToolParameter>, |
1115 | 35 | ) -> Self { |
1116 | 35 | Self { |
1117 | 35 | name: name.into(), |
1118 | 35 | description: description.into(), |
1119 | 35 | parameters, |
1120 | 35 | } |
1121 | 35 | } |
1122 | | |
1123 | | /// Get required parameters |
1124 | 1 | pub fn required_params(&self) -> impl Iterator<Item = &ToolParameter> { |
1125 | 1 | self.parameters.iter().filter(|p| p.required) |
1126 | 1 | } |
1127 | | |
1128 | | /// Get optional parameters |
1129 | 1 | pub fn optional_params(&self) -> impl Iterator<Item = &ToolParameter> { |
1130 | 3 | self.parameters.iter()1 .filter1 (|p| !p.required) |
1131 | 1 | } |
1132 | | |
1133 | | /// Validate tool name is a valid identifier |
1134 | 10 | pub fn is_valid_name(name: &str) -> bool { |
1135 | 10 | if name.is_empty() { |
1136 | 1 | return false; |
1137 | 9 | } |
1138 | 9 | let mut chars = name.chars(); |
1139 | | // SAFETY: name is non-empty (checked above) |
1140 | 9 | let first = chars.next().expect("name is non-empty"); |
1141 | 9 | if !first.is_ascii_alphabetic() && first != '_'2 { |
1142 | 1 | return false; |
1143 | 8 | } |
1144 | 52 | chars8 .all8 (|c| c.is_ascii_alphanumeric() || c == '_'4 ) |
1145 | 10 | } |
1146 | | } |
1147 | | |
1148 | | /// Tool choice mode |
1149 | | #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] |
1150 | | #[serde(rename_all = "snake_case")] |
1151 | | pub enum ToolChoice { |
1152 | | /// Model decides whether to call a tool |
1153 | | #[default] |
1154 | | Auto, |
1155 | | /// Model must call at least one tool |
1156 | | Required, |
1157 | | /// Model must not call any tools |
1158 | | None, |
1159 | | /// Model must call the specified tool |
1160 | | Specific(String), |
1161 | | } |
1162 | | |
1163 | | /// Parsed tool call from model output |
1164 | | #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] |
1165 | | pub struct ToolCall { |
1166 | | /// Unique ID for this tool call |
1167 | | pub id: String, |
1168 | | /// Tool name |
1169 | | pub name: String, |
1170 | | /// Tool arguments as JSON string |
1171 | | pub arguments: String, |
1172 | | } |
1173 | | |
1174 | | impl ToolCall { |
1175 | | /// Create new tool call |
1176 | 17 | pub fn new( |
1177 | 17 | id: impl Into<String>, |
1178 | 17 | name: impl Into<String>, |
1179 | 17 | arguments: impl Into<String>, |
1180 | 17 | ) -> Self { |
1181 | 17 | Self { |
1182 | 17 | id: id.into(), |
1183 | 17 | name: name.into(), |
1184 | 17 | arguments: arguments.into(), |
1185 | 17 | } |
1186 | 17 | } |
1187 | | |
1188 | | /// Parse arguments as JSON value |
1189 | | /// |
1190 | | /// # Errors |
1191 | | /// |
1192 | | /// Returns `RealizarError::InvalidConfiguration` if the arguments are not valid JSON. |
1193 | 2 | pub fn parse_arguments(&self) -> Result<serde_json::Value> { |
1194 | 2 | serde_json::from_str(&self.arguments).map_err(|e| {1 |
1195 | 1 | RealizarError::InvalidConfiguration(format!("Failed to parse tool arguments: {e}")) |
1196 | 1 | }) |
1197 | 2 | } |
1198 | | } |
1199 | | |
1200 | | /// Tool call result |
1201 | | #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] |
1202 | | pub struct ToolResult { |
1203 | | /// Tool call ID this result is for |
1204 | | pub tool_call_id: String, |
1205 | | /// Result content (JSON or plain text) |
1206 | | pub content: String, |
1207 | | /// Whether the tool call succeeded |
1208 | | #[serde(default = "default_true")] |
1209 | | pub success: bool, |
1210 | | } |
1211 | | |
1212 | 1 | fn default_true() -> bool { |
1213 | 1 | true |
1214 | 1 | } |
1215 | | |
1216 | | impl ToolResult { |
1217 | | /// Create successful result |
1218 | 2 | pub fn success(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self { |
1219 | 2 | Self { |
1220 | 2 | tool_call_id: tool_call_id.into(), |
1221 | 2 | content: content.into(), |
1222 | 2 | success: true, |
1223 | 2 | } |
1224 | 2 | } |
1225 | | |
1226 | | /// Create error result |
1227 | 1 | pub fn error(tool_call_id: impl Into<String>, error: impl Into<String>) -> Self { |
1228 | 1 | Self { |
1229 | 1 | tool_call_id: tool_call_id.into(), |
1230 | 1 | content: error.into(), |
1231 | 1 | success: false, |
1232 | 1 | } |
1233 | 1 | } |
1234 | | } |
1235 | | |
1236 | | /// Tool call parser for extracting tool calls from model output |
1237 | | pub struct ToolCallParser { |
1238 | | /// Available tools |
1239 | | tools: Vec<ToolDefinition>, |
1240 | | /// Tool call format (default: OpenAI-style JSON) |
1241 | | format: ToolCallFormat, |
1242 | | /// Next tool call ID |
1243 | | next_id: u64, |
1244 | | } |
1245 | | |
1246 | | /// Format for tool calls in model output |
1247 | | #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] |
1248 | | pub enum ToolCallFormat { |
1249 | | /// OpenAI-style JSON: {"name": "tool", "arguments": {...}} |
1250 | | #[default] |
1251 | | OpenAI, |
1252 | | /// Anthropic-style XML: <tool_use><name>tool</name><input>{...}</input></tool_use> |
1253 | | Anthropic, |
1254 | | /// Hermes format: <tool_call>{"name": "tool", "arguments": {...}}</tool_call> |
1255 | | Hermes, |
1256 | | } |
1257 | | |
1258 | | impl ToolCallParser { |
1259 | | /// Create new parser with tools |
1260 | 23 | pub fn new(tools: Vec<ToolDefinition>) -> Self { |
1261 | 23 | Self { |
1262 | 23 | tools, |
1263 | 23 | format: ToolCallFormat::default(), |
1264 | 23 | next_id: 0, |
1265 | 23 | } |
1266 | 23 | } |
1267 | | |
1268 | | /// Set tool call format |
1269 | | #[must_use] |
1270 | 12 | pub fn with_format(mut self, format: ToolCallFormat) -> Self { |
1271 | 12 | self.format = format; |
1272 | 12 | self |
1273 | 12 | } |
1274 | | |
1275 | | /// Generate unique tool call ID |
1276 | 16 | pub fn generate_id(&mut self) -> String { |
1277 | 16 | let id = format!("call_{}", self.next_id); |
1278 | 16 | self.next_id += 1; |
1279 | 16 | id |
1280 | 16 | } |
1281 | | |
1282 | | /// Get available tool names |
1283 | 2 | pub fn tool_names(&self) -> impl Iterator<Item = &str> { |
1284 | 2 | self.tools.iter().map(|t| t.name.as_str()) |
1285 | 2 | } |
1286 | | |
1287 | | /// Get tool by name |
1288 | 3 | pub fn get_tool(&self, name: &str) -> Option<&ToolDefinition> { |
1289 | 5 | self.tools.iter()3 .find3 (|t| t.name == name) |
1290 | 3 | } |
1291 | | |
1292 | | /// Parse tool calls from text (OpenAI format) |
1293 | | /// |
1294 | | /// Looks for JSON objects with "name" and "arguments" fields. |
1295 | 19 | pub fn parse(&mut self, text: &str) -> Vec<ToolCall> { |
1296 | 19 | match self.format { |
1297 | 8 | ToolCallFormat::OpenAI => self.parse_openai(text), |
1298 | 5 | ToolCallFormat::Anthropic => self.parse_anthropic(text), |
1299 | 6 | ToolCallFormat::Hermes => self.parse_hermes(text), |
1300 | | } |
1301 | 19 | } |
1302 | | |
1303 | 8 | fn parse_openai(&mut self, text: &str) -> Vec<ToolCall> { |
1304 | 8 | let mut calls = Vec::new(); |
1305 | | |
1306 | | // Try to find JSON objects with tool call structure |
1307 | | // Look for patterns like {"name": "...", "arguments": {...}} |
1308 | 8 | let mut start = 0; |
1309 | 19 | while let Some(pos11 ) = text[start..].find('{') { |
1310 | 11 | let abs_pos = start + pos; |
1311 | 11 | if let Some(end10 ) = find_matching_brace(&text[abs_pos..]) { |
1312 | 10 | let json_str = &text[abs_pos..=(abs_pos + end)]; |
1313 | 10 | if let Ok(value9 ) = serde_json::from_str::<serde_json::Value>(json_str) { |
1314 | 7 | if let (Some(name), Some(args)) = ( |
1315 | 9 | value.get("name").and_then(|v| v7 .as_str7 ()), |
1316 | 9 | value.get("arguments"), |
1317 | | ) { |
1318 | | // Check if this is a valid tool |
1319 | 8 | if self.tools.iter()7 .any7 (|t| t.name == name) { |
1320 | 6 | let arguments = if args.is_string() { |
1321 | 1 | args.as_str().expect("checked is_string above").to_string() |
1322 | | } else { |
1323 | 5 | args.to_string() |
1324 | | }; |
1325 | 6 | calls.push(ToolCall::new(self.generate_id(), name, arguments)); |
1326 | 1 | } |
1327 | 2 | } |
1328 | 1 | } |
1329 | 10 | start = abs_pos + end + 1; |
1330 | 1 | } else { |
1331 | 1 | start = abs_pos + 1; |
1332 | 1 | } |
1333 | | } |
1334 | | |
1335 | 8 | calls |
1336 | 8 | } |
1337 | | |
1338 | 5 | fn parse_anthropic(&mut self, text: &str) -> Vec<ToolCall> { |
1339 | 5 | let mut calls = Vec::new(); |
1340 | | |
1341 | | // Look for <tool_use>...</tool_use> blocks |
1342 | 5 | let mut pos = 0; |
1343 | 10 | while let Some(start6 ) = text[pos..].find("<tool_use>") { |
1344 | 6 | let abs_start = pos + start + 10; // Skip "<tool_use>" |
1345 | 6 | if let Some(end5 ) = text[abs_start..].find("</tool_use>") { |
1346 | 5 | let content = &text[abs_start..abs_start + end]; |
1347 | | |
1348 | | // Extract name |
1349 | 5 | let name = extract_xml_tag(content, "name"); |
1350 | 5 | let input = extract_xml_tag(content, "input"); |
1351 | | |
1352 | 5 | if let (Some(name4 ), Some(input4 )) = (name, input) { |
1353 | 5 | if self.tools.iter()4 .any4 (|t| t.name == name) { |
1354 | 3 | calls.push(ToolCall::new(self.generate_id(), name, input)); |
1355 | 3 | }1 |
1356 | 1 | } |
1357 | | |
1358 | 5 | pos = abs_start + end + 11; // Skip "</tool_use>" |
1359 | | } else { |
1360 | 1 | break; |
1361 | | } |
1362 | | } |
1363 | | |
1364 | 5 | calls |
1365 | 5 | } |
1366 | | |
1367 | 6 | fn parse_hermes(&mut self, text: &str) -> Vec<ToolCall> { |
1368 | 6 | let mut calls = Vec::new(); |
1369 | | |
1370 | | // Look for <tool_call>...</tool_call> blocks |
1371 | 6 | let mut pos = 0; |
1372 | 12 | while let Some(start7 ) = text[pos..].find("<tool_call>") { |
1373 | 7 | let abs_start = pos + start + 11; // Skip "<tool_call>" |
1374 | 7 | if let Some(end6 ) = text[abs_start..].find("</tool_call>") { |
1375 | 6 | let json_str = text[abs_start..abs_start + end].trim(); |
1376 | | |
1377 | 6 | if let Ok(value5 ) = serde_json::from_str::<serde_json::Value>(json_str) { |
1378 | 5 | if let (Some(name), Some(args)) = ( |
1379 | 5 | value.get("name").and_then(|v| v.as_str()), |
1380 | 5 | value.get("arguments"), |
1381 | | ) { |
1382 | 6 | if self.tools.iter()5 .any5 (|t| t.name == name) { |
1383 | 4 | let arguments = if args.is_string() { |
1384 | 1 | args.as_str().expect("checked is_string above").to_string() |
1385 | | } else { |
1386 | 3 | args.to_string() |
1387 | | }; |
1388 | 4 | calls.push(ToolCall::new(self.generate_id(), name, arguments)); |
1389 | 1 | } |
1390 | 0 | } |
1391 | 1 | } |
1392 | | |
1393 | 6 | pos = abs_start + end + 12; // Skip "</tool_call>" |
1394 | | } else { |
1395 | 1 | break; |
1396 | | } |
1397 | | } |
1398 | | |
1399 | 6 | calls |
1400 | 6 | } |
1401 | | } |
1402 | | |
1403 | | /// Find matching closing brace, handling nested braces |
1404 | 18 | fn find_matching_brace(s: &str) -> Option<usize> { |
1405 | 18 | let mut depth = 0; |
1406 | 18 | let mut in_string = false; |
1407 | 18 | let mut escape_next = false; |
1408 | | |
1409 | 648 | for (i, c) in s18 .char_indices18 () { |
1410 | 648 | if escape_next { |
1411 | 7 | escape_next = false; |
1412 | 7 | continue; |
1413 | 641 | } |
1414 | | |
1415 | 7 | match c { |
1416 | 7 | '\\' if in_string => escape_next = true, |
1417 | 104 | '"' => in_string = !in_string, |
1418 | 27 | '{' if !in_string => depth += 1, |
1419 | 25 | '}' if !in_string => { |
1420 | 25 | depth -= 1; |
1421 | 25 | if depth == 0 { |
1422 | 15 | return Some(i); |
1423 | 10 | } |
1424 | | }, |
1425 | 478 | _ => {}, |
1426 | | } |
1427 | | } |
1428 | | |
1429 | 3 | None |
1430 | 18 | } |
1431 | | |
1432 | | /// Extract content of an XML tag |
1433 | 14 | fn extract_xml_tag(content: &str, tag: &str) -> Option<String> { |
1434 | 14 | let open_tag = format!("<{tag}>"); |
1435 | 14 | let close_tag = format!("</{tag}>"); |
1436 | | |
1437 | 14 | if let Some(start12 ) = content.find(&open_tag) { |
1438 | 12 | let value_start = start + open_tag.len(); |
1439 | 12 | if let Some(end11 ) = content[value_start..].find(&close_tag) { |
1440 | 11 | return Some(content[value_start..value_start + end].to_string()); |
1441 | 1 | } |
1442 | 2 | } |
1443 | 3 | None |
1444 | 14 | } |
1445 | | |
1446 | | /// Generate grammar for tool calling output |
1447 | | /// |
1448 | | /// Creates a grammar that constrains model output to valid tool calls. |
1449 | 5 | pub fn generate_tool_grammar(tools: &[ToolDefinition]) -> Grammar { |
1450 | 5 | let mut grammar = Grammar::default(); |
1451 | | |
1452 | | // Add basic rules |
1453 | 5 | add_json_whitespace_rules(&mut grammar); |
1454 | | |
1455 | | // Generate alternatives for each tool |
1456 | 5 | let mut tool_alternatives = Vec::new(); |
1457 | | |
1458 | 10 | for tool5 in tools { |
1459 | 5 | let tool_rule = format!("tool_{}", tool.name); |
1460 | | |
1461 | | // Generate parameter object grammar |
1462 | 5 | let params_rule = format!("{tool_rule}_params"); |
1463 | 5 | generate_params_grammar(&mut grammar, ¶ms_rule, &tool.parameters); |
1464 | | |
1465 | | // Tool call rule: {"name": "tool_name", "arguments": {...}} |
1466 | 5 | let mut elements = vec![ |
1467 | 5 | GrammarElement::Char('{'), |
1468 | 5 | GrammarElement::RuleRef("ws".to_string()), |
1469 | 5 | GrammarElement::Char('"'), |
1470 | 5 | GrammarElement::Char('n'), |
1471 | 5 | GrammarElement::Char('a'), |
1472 | 5 | GrammarElement::Char('m'), |
1473 | 5 | GrammarElement::Char('e'), |
1474 | 5 | GrammarElement::Char('"'), |
1475 | 5 | GrammarElement::RuleRef("ws".to_string()), |
1476 | 5 | GrammarElement::Char(':'), |
1477 | 5 | GrammarElement::RuleRef("ws".to_string()), |
1478 | 5 | GrammarElement::Char('"'), |
1479 | | ]; |
1480 | | |
1481 | | // Tool name literal |
1482 | 49 | for c in tool.name5 .chars5 () { |
1483 | 49 | elements.push(GrammarElement::Char(c)); |
1484 | 49 | } |
1485 | | |
1486 | 5 | elements.extend([ |
1487 | 5 | GrammarElement::Char('"'), |
1488 | 5 | GrammarElement::RuleRef("ws".to_string()), |
1489 | 5 | GrammarElement::Char(','), |
1490 | 5 | GrammarElement::RuleRef("ws".to_string()), |
1491 | 5 | GrammarElement::Char('"'), |
1492 | 5 | GrammarElement::Char('a'), |
1493 | 5 | GrammarElement::Char('r'), |
1494 | 5 | GrammarElement::Char('g'), |
1495 | 5 | GrammarElement::Char('u'), |
1496 | 5 | GrammarElement::Char('m'), |
1497 | 5 | GrammarElement::Char('e'), |
1498 | 5 | GrammarElement::Char('n'), |
1499 | 5 | GrammarElement::Char('t'), |
1500 | 5 | GrammarElement::Char('s'), |
1501 | 5 | GrammarElement::Char('"'), |
1502 | 5 | GrammarElement::RuleRef("ws".to_string()), |
1503 | 5 | GrammarElement::Char(':'), |
1504 | 5 | GrammarElement::RuleRef("ws".to_string()), |
1505 | 5 | GrammarElement::RuleRef(params_rule), |
1506 | 5 | GrammarElement::RuleRef("ws".to_string()), |
1507 | 5 | GrammarElement::Char('}'), |
1508 | 5 | ]); |
1509 | | |
1510 | 5 | grammar.add_rule(GrammarRule::new( |
1511 | 5 | &tool_rule, |
1512 | 5 | vec![GrammarAlternative::new(elements)], |
1513 | | )); |
1514 | | |
1515 | 5 | tool_alternatives.push(GrammarAlternative::new(vec![GrammarElement::RuleRef( |
1516 | 5 | tool_rule, |
1517 | 5 | )])); |
1518 | | } |
1519 | | |
1520 | | // Root rule: one of the tools |
1521 | 5 | grammar.add_rule(GrammarRule::new("root", tool_alternatives)); |
1522 | | |
1523 | 5 | grammar |
1524 | 5 | } |
1525 | | |
1526 | | /// Generate grammar for tool parameters |
1527 | 6 | fn generate_params_grammar(grammar: &mut Grammar, rule_name: &str, params: &[ToolParameter]) { |
1528 | 6 | if params.is_empty() { |
1529 | | // Empty object: {} |
1530 | 3 | grammar.add_rule(GrammarRule::new( |
1531 | 3 | rule_name, |
1532 | 3 | vec![GrammarAlternative::new(vec![ |
1533 | 3 | GrammarElement::Char('{'), |
1534 | 3 | GrammarElement::RuleRef("ws".to_string()), |
1535 | 3 | GrammarElement::Char('}'), |
1536 | | ])], |
1537 | | )); |
1538 | 3 | return; |
1539 | 3 | } |
1540 | | |
1541 | | // Build object with all parameters |
1542 | 3 | let mut elements = vec![ |
1543 | 3 | GrammarElement::Char('{'), |
1544 | 3 | GrammarElement::RuleRef("ws".to_string()), |
1545 | | ]; |
1546 | | |
1547 | 9 | for (i, param) in params3 .iter3 ().enumerate3 () { |
1548 | 9 | if i > 0 { |
1549 | 6 | elements.push(GrammarElement::Char(',')); |
1550 | 6 | elements.push(GrammarElement::RuleRef("ws".to_string())); |
1551 | 6 | }3 |
1552 | | |
1553 | | // Property name |
1554 | 9 | elements.push(GrammarElement::Char('"')); |
1555 | 84 | for c in param.name9 .chars9 () { |
1556 | 84 | elements.push(GrammarElement::Char(c)); |
1557 | 84 | } |
1558 | 9 | elements.push(GrammarElement::Char('"')); |
1559 | 9 | elements.push(GrammarElement::RuleRef("ws".to_string())); |
1560 | 9 | elements.push(GrammarElement::Char(':')); |
1561 | 9 | elements.push(GrammarElement::RuleRef("ws".to_string())); |
1562 | | |
1563 | | // Property value based on type |
1564 | 9 | let value_rule = format!("{rule_name}_{}", param.name); |
1565 | 9 | generate_param_type_grammar(grammar, &value_rule, ¶m.param_type); |
1566 | 9 | elements.push(GrammarElement::RuleRef(value_rule)); |
1567 | | } |
1568 | | |
1569 | 3 | elements.push(GrammarElement::RuleRef("ws".to_string())); |
1570 | 3 | elements.push(GrammarElement::Char('}')); |
1571 | | |
1572 | 3 | grammar.add_rule(GrammarRule::new( |
1573 | 3 | rule_name, |
1574 | 3 | vec![GrammarAlternative::new(elements)], |
1575 | | )); |
1576 | 6 | } |
1577 | | |
1578 | | /// Generate grammar for a parameter type |
1579 | 10 | fn generate_param_type_grammar( |
1580 | 10 | grammar: &mut Grammar, |
1581 | 10 | rule_name: &str, |
1582 | 10 | param_type: &ToolParameterType, |
1583 | 10 | ) { |
1584 | 10 | match param_type { |
1585 | 3 | ToolParameterType::String => { |
1586 | 3 | grammar.add_rule(GrammarRule::new( |
1587 | 3 | rule_name, |
1588 | 3 | vec![GrammarAlternative::new(vec![GrammarElement::RuleRef( |
1589 | 3 | "string".to_string(), |
1590 | 3 | )])], |
1591 | 3 | )); |
1592 | 3 | }, |
1593 | 2 | ToolParameterType::Integer => { |
1594 | 2 | grammar.add_rule(GrammarRule::new( |
1595 | 2 | rule_name, |
1596 | 2 | vec![GrammarAlternative::new(vec![GrammarElement::RuleRef( |
1597 | 2 | "integer".to_string(), |
1598 | 2 | )])], |
1599 | 2 | )); |
1600 | 2 | }, |
1601 | 1 | ToolParameterType::Number => { |
1602 | 1 | grammar.add_rule(GrammarRule::new( |
1603 | 1 | rule_name, |
1604 | 1 | vec![GrammarAlternative::new(vec![GrammarElement::RuleRef( |
1605 | 1 | "number".to_string(), |
1606 | 1 | )])], |
1607 | 1 | )); |
1608 | 1 | }, |
1609 | 1 | ToolParameterType::Boolean => { |
1610 | 1 | grammar.add_rule(GrammarRule::new( |
1611 | 1 | rule_name, |
1612 | 1 | vec![GrammarAlternative::new(vec![GrammarElement::RuleRef( |
1613 | 1 | "boolean".to_string(), |
1614 | 1 | )])], |
1615 | 1 | )); |
1616 | 1 | }, |
1617 | 1 | ToolParameterType::Enum(values) => { |
1618 | 1 | let alternatives: Vec<_> = values |
1619 | 1 | .iter() |
1620 | 2 | .map1 (|v| { |
1621 | 2 | let mut chars = vec![GrammarElement::Char('"')]; |
1622 | 2 | chars.extend(v.chars().map(GrammarElement::Char)); |
1623 | 2 | chars.push(GrammarElement::Char('"')); |
1624 | 2 | GrammarAlternative::new(chars) |
1625 | 2 | }) |
1626 | 1 | .collect(); |
1627 | 1 | grammar.add_rule(GrammarRule::new(rule_name, alternatives)); |
1628 | | }, |
1629 | 1 | ToolParameterType::Array { items } => { |
1630 | 1 | let item_rule = format!("{rule_name}_item"); |
1631 | 1 | generate_param_type_grammar(grammar, &item_rule, items); |
1632 | 1 | |
1633 | 1 | let items_rule = format!("{rule_name}_items"); |
1634 | 1 | grammar.add_rule(GrammarRule::new( |
1635 | 1 | &items_rule, |
1636 | 1 | vec![ |
1637 | 1 | GrammarAlternative::new(vec![]), // Empty |
1638 | 1 | GrammarAlternative::new(vec![ |
1639 | 1 | GrammarElement::Char(','), |
1640 | 1 | GrammarElement::RuleRef("ws".to_string()), |
1641 | 1 | GrammarElement::RuleRef(item_rule.clone()), |
1642 | 1 | GrammarElement::RuleRef(items_rule.clone()), |
1643 | 1 | ]), |
1644 | 1 | ], |
1645 | 1 | )); |
1646 | 1 | |
1647 | 1 | grammar.add_rule(GrammarRule::new( |
1648 | 1 | rule_name, |
1649 | 1 | vec![ |
1650 | 1 | GrammarAlternative::new(vec![ |
1651 | 1 | GrammarElement::Char('['), |
1652 | 1 | GrammarElement::RuleRef("ws".to_string()), |
1653 | 1 | GrammarElement::Char(']'), |
1654 | 1 | ]), |
1655 | 1 | GrammarAlternative::new(vec![ |
1656 | 1 | GrammarElement::Char('['), |
1657 | 1 | GrammarElement::RuleRef("ws".to_string()), |
1658 | 1 | GrammarElement::RuleRef(item_rule), |
1659 | 1 | GrammarElement::RuleRef(items_rule), |
1660 | 1 | GrammarElement::RuleRef("ws".to_string()), |
1661 | 1 | GrammarElement::Char(']'), |
1662 | 1 | ]), |
1663 | 1 | ], |
1664 | 1 | )); |
1665 | 1 | }, |
1666 | 1 | ToolParameterType::Object { properties } => { |
1667 | 1 | generate_params_grammar(grammar, rule_name, properties); |
1668 | 1 | }, |
1669 | | } |
1670 | 10 | } |
1671 | | |
1672 | | /// Add standard JSON whitespace rules to grammar |
1673 | 5 | fn add_json_whitespace_rules(grammar: &mut Grammar) { |
1674 | | // Whitespace (optional) |
1675 | 5 | grammar.add_rule(GrammarRule::new( |
1676 | | "ws", |
1677 | 5 | vec![ |
1678 | 5 | GrammarAlternative::new(vec![]), |
1679 | 5 | GrammarAlternative::new(vec![ |
1680 | 5 | GrammarElement::Char(' '), |
1681 | 5 | GrammarElement::RuleRef("ws".to_string()), |
1682 | | ]), |
1683 | 5 | GrammarAlternative::new(vec![ |
1684 | 5 | GrammarElement::Char('\n'), |
1685 | 5 | GrammarElement::RuleRef("ws".to_string()), |
1686 | | ]), |
1687 | 5 | GrammarAlternative::new(vec![ |
1688 | 5 | GrammarElement::Char('\t'), |
1689 | 5 | GrammarElement::RuleRef("ws".to_string()), |
1690 | | ]), |
1691 | | ], |
1692 | | )); |
1693 | | |
1694 | | // String (simplified - no escape handling) |
1695 | 5 | grammar.add_rule(GrammarRule::new( |
1696 | | "string", |
1697 | 5 | vec![GrammarAlternative::new(vec![ |
1698 | 5 | GrammarElement::Char('"'), |
1699 | 5 | GrammarElement::RuleRef("string_chars".to_string()), |
1700 | 5 | GrammarElement::Char('"'), |
1701 | | ])], |
1702 | | )); |
1703 | | |
1704 | 5 | grammar.add_rule(GrammarRule::new( |
1705 | | "string_chars", |
1706 | 5 | vec![ |
1707 | 5 | GrammarAlternative::new(vec![]), |
1708 | 5 | GrammarAlternative::new(vec![ |
1709 | 5 | GrammarElement::CharNot(vec!['"', '\\']), |
1710 | 5 | GrammarElement::RuleRef("string_chars".to_string()), |
1711 | | ]), |
1712 | | ], |
1713 | | )); |
1714 | | |
1715 | | // Integer |
1716 | 5 | grammar.add_rule(GrammarRule::new( |
1717 | | "integer", |
1718 | 5 | vec![ |
1719 | 5 | GrammarAlternative::new(vec![ |
1720 | 5 | GrammarElement::Char('-'), |
1721 | 5 | GrammarElement::RuleRef("digits".to_string()), |
1722 | | ]), |
1723 | 5 | GrammarAlternative::new(vec![GrammarElement::RuleRef("digits".to_string())]), |
1724 | | ], |
1725 | | )); |
1726 | | |
1727 | | // Number (with optional decimal) |
1728 | 5 | grammar.add_rule(GrammarRule::new( |
1729 | | "number", |
1730 | 5 | vec![ |
1731 | 5 | GrammarAlternative::new(vec![ |
1732 | 5 | GrammarElement::RuleRef("integer".to_string()), |
1733 | 5 | GrammarElement::Char('.'), |
1734 | 5 | GrammarElement::RuleRef("digits".to_string()), |
1735 | | ]), |
1736 | 5 | GrammarAlternative::new(vec![GrammarElement::RuleRef("integer".to_string())]), |
1737 | | ], |
1738 | | )); |
1739 | | |
1740 | | // Digits |
1741 | 5 | grammar.add_rule(GrammarRule::new( |
1742 | | "digits", |
1743 | 5 | vec![ |
1744 | 5 | GrammarAlternative::new(vec![ |
1745 | 5 | GrammarElement::CharRange('0', '9'), |
1746 | 5 | GrammarElement::RuleRef("digits".to_string()), |
1747 | | ]), |
1748 | 5 | GrammarAlternative::new(vec![GrammarElement::CharRange('0', '9')]), |
1749 | | ], |
1750 | | )); |
1751 | | |
1752 | | // Boolean |
1753 | 5 | grammar.add_rule(GrammarRule::new( |
1754 | | "boolean", |
1755 | 5 | vec![ |
1756 | 5 | GrammarAlternative::new(vec![ |
1757 | 5 | GrammarElement::Char('t'), |
1758 | 5 | GrammarElement::Char('r'), |
1759 | 5 | GrammarElement::Char('u'), |
1760 | 5 | GrammarElement::Char('e'), |
1761 | | ]), |
1762 | 5 | GrammarAlternative::new(vec![ |
1763 | 5 | GrammarElement::Char('f'), |
1764 | 5 | GrammarElement::Char('a'), |
1765 | 5 | GrammarElement::Char('l'), |
1766 | 5 | GrammarElement::Char('s'), |
1767 | 5 | GrammarElement::Char('e'), |
1768 | | ]), |
1769 | | ], |
1770 | | )); |
1771 | 5 | } |
1772 | | |
1773 | | // ============================================================================= |
1774 | | // TESTS |
1775 | | // ============================================================================= |
1776 | | |
1777 | | // Tests extracted to tests.rs (PMAT-802) |
1778 | | #[cfg(test)] |
1779 | | #[path = "tests.rs"] |
1780 | | mod grammar_tests; |