aufbau 0.1.0

Type-aware constrained decoding for LLMs using context-dependent grammars with typing rules
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86

use crate::logic::grammar::Grammar;
use crate::logic::partial::Synthesizer;
use crate::logic::typing::Context;

const SCOPED_TYPED_SPEC: &str = r#"
    Identifier ::= /[a-z]+/
    Type ::= 'X' | 'Y'
    Variable(var) ::= Identifier[x]
    Num(num) ::= /[0-9]+/
    Let(letb) ::= 'def' Identifier[name] ':' Type[τ] '=' Atom[value] 'in' Expr[body]
    Scoped(scoped) ::= '{' Expr[inner] '}'
    Atom ::= Variable | Num | Scoped | '(' Expr ')'
    Expr ::= Let | Atom

    x ∈ Γ
    ----------- (var)
    Γ(x)

    ----------- (num)
    'X'

    Γ ⊢ value : τ, Γ[name:τ] ⊢ body : ?R
    ----------- (letb)
    ?R

    [Γ] ⊢ inner : ?T
    ----------- (scoped)
    ?T
"#;

/// "def a : X =" should have typed completions (the value position expects an Atom).
#[test]
fn typed_completions_for_def_after_equals() {
    let grammar = Grammar::load(SCOPED_TYPED_SPEC).unwrap();
    let ctx = Context::new();
    let mut synth = Synthesizer::new(grammar.clone(), "def a : X =");
    let tokens = synth.completions_ctx(&ctx);
    assert!(
        !tokens.is_empty(),
        "should have typed completions for 'def a : X ='"
    );
}

/// try_extend from "def a : X =" with a numeric literal must succeed.
#[test]
fn try_extend_def_after_equals_with_number() {
    let grammar = Grammar::load(SCOPED_TYPED_SPEC).unwrap();
    let ctx = Context::new();
    let mut synth = Synthesizer::new(grammar.clone(), "def a : X =");
    assert!(
        synth.try_extend("1", &ctx).is_ok(),
        "should extend 'def a : X =' with '1'"
    );
}

#[test]
fn extend_commits_input_and_tree() {
    let grammar = Grammar::load(SCOPED_TYPED_SPEC).unwrap();
    let ctx = Context::new();
    let mut synth = Synthesizer::new(grammar, "def a : X =");

    let typed = synth.extend("1", &ctx).expect("extend should succeed");

    assert_eq!(synth.input(), "def a : X = 1");
    assert_eq!(synth.tree().unwrap().text(), typed.text());
}

#[test]
fn extend_with_regex_commits_input_and_tree() {
    let grammar = Grammar::load(SCOPED_TYPED_SPEC).unwrap();
    let ctx = Context::new();
    let mut synth = Synthesizer::new(grammar, "def a : X =");
    let token = synth
        .completions_ctx(&ctx)
        .iter()
        .next()
        .cloned()
        .expect("expected at least one completion token");

    let (typed, extended) = synth
        .extend_with_regex(&token, &ctx)
        .expect("regex extension should succeed");

    assert_eq!(synth.input(), extended);
    assert_eq!(synth.tree().unwrap().text(), typed.text());
}