Swift と GBNF 文法でローカル LLM に必ず正しい JSON を出力させる

struct Letter: Codable {
    var sender: String
    var action: String      // "none", "pay", "reply" or "sign"
    var urgency: Int        // 0...3
    var deadline: String    // "YYYY-MM-DD" or ""
    var steps: [String]     // at most 4
}

{"sender": "City Energy", "action": "pay", "urgency": 2, "deadline": "2026-10-30",
 "steps": ["Open your banking app", "Transfer 84.20 EUR", "Use reference 4711"]}

root ::= "{" ws "\"sender\":" ws str "," ws "\"action\":" ws action "," ws "\"urgency\":" ws urgency "," ws "\"deadline\":" ws date "," ws "\"steps\":" ws steps ws "}"
action ::= "\"none\"" | "\"pay\"" | "\"reply\"" | "\"sign\""
urgency ::= [0-3]
steps ::= "[" ws "]" | "[" ws str ( ws "," ws str ){0,3} ws "]"
str ::= "\"" chr* "\""
chr ::= [^"\\\x00-\x1f] | "\\" ["\\/bfnrt] | "\\u" [0-9a-fA-F]{4}
date ::= "\"\"" | "\"" [0-9]{4} "-" [0-9]{2} "-" [0-9]{2} "\""
ws ::= [ \t\n]{0,8}

indirect enum Field {
    case string
    case choice([String])
    case digit(ClosedRange<Int>)       // a single digit, e.g. 0...3
    case date
    case list(Field, max: Int)
}

struct Key {
    let name: String                   // letters and digits only: it doubles as a rule name
    let field: Field
    let hint: String                   // goes into the prompt, not the grammar
}

func gbnf(_ keys: [Key]) -> String {
    var rules: [String] = []

    func rule(_ field: Field, _ name: String) -> String {
        switch field {
        case .string: return "str"
        case .date: return "date"
        case .digit(let r):
            rules.append("\(name) ::= [\(r.lowerBound)-\(r.upperBound)]")
            return name
        case .choice(let values):
            let alternatives = values.map { #""\"\#($0)\"""# }.joined(separator: " | ")
            rules.append("\(name) ::= \(alternatives)")
            return name
        case .list(let item, let max):
            let i = rule(item, name + "-item")
            rules.append(#"\#(name) ::= "[" ws "]" | "[" ws \#(i) ( ws "," ws \#(i) ){0,\#(max - 1)} ws "]""#)
            return name
        }
    }

    let fields = keys.map { k in #""\"\#(k.name)\":" ws "# + rule(k.field, k.name) }
    let root = #"root ::= "{" ws "# + fields.joined(separator: #" "," ws "#) + #" ws "}""#
    let shared = [
        #"str ::= "\"" chr* "\"""#,
        #"chr ::= [^"\\\x00-\x1f] | "\\" ["\\/bfnrt] | "\\u" [0-9a-fA-F]{4}"#,
        #"date ::= "\"\"" | "\"" [0-9]{4} "-" [0-9]{2} "-" [0-9]{2} "\"""#,
        #"ws ::= [ \t\n]{0,8}"#,
    ]
    return ([root] + rules + shared).joined(separator: "\n")
}

let letterKeys = [
    Key(name: "sender", field: .string, hint: "Who sent the letter, short"),
    Key(name: "action", field: .choice(["none", "pay", "reply", "sign"]), hint: "What the reader has to do"),
    Key(name: "urgency", field: .digit(0...3), hint: "0 only information, 3 deadline soon"),
    Key(name: "deadline", field: .date, hint: "Only if the letter names one"),
    Key(name: "steps", field: .list(.string, max: 4), hint: "Tiny steps, each starting with a verb"),
]

func describe(_ keys: [Key]) -> String {
    keys.map { "- \($0.name): \($0.hint)" }.joined(separator: "\n")
}

let system = """
    You help people deal with paper mail. Use only facts from the letter.
    Answer only with JSON with these keys:
    \(describe(letterKeys))
    """

guard let sampler = llama_sampler_chain_init(llama_sampler_chain_default_params()) else { throw LLMError.decodeFailed }
defer { llama_sampler_free(sampler) }

guard let grammar = llama_sampler_init_grammar(vocab, gbnf(letterKeys), "root") else {
    throw LLMError.decodeFailed            // the grammar text does not parse
}
llama_sampler_chain_add(sampler, grammar)
llama_sampler_chain_add(sampler, llama_sampler_init_greedy())

let raw = try await llm.generate(messages, grammar: gbnf(letterKeys), maxTokens: 1024)
let letter = try JSONDecoder().decode(Letter.self, from: Data(raw.utf8))