Widen the empty-array probe to merged item-start tokens

Sources/AnyLanguageModel/Shared/StructuredGeneration.swift

@@ -580,39 +580,80 @@ struct ConstrainedJSONGenerator<Backend: TokenBackend> {
 
     /// Probe after `[` when `minItems == 0`: model may close immediately or start an item.
     ///
-    /// Samples once among `]` and tokens that can start the item type. Choosing `]` means
-    /// close; any other sample is discarded without decoding so ``generateNode`` can emit
-    /// the first element from the same backend state.
+    /// Samples once among closing tokens and tokens that can start the item type. Choosing
+    /// a close means close; any other sample is discarded without decoding so
+    /// ``generateNode`` can emit the first element from the same backend state.
+    ///
+    /// Byte-pair vocabularies put most item-start probability on merged tokens (a quote
+    /// fused with the first word, or a leading space), so the probe admits every token
+    /// whose first non-whitespace character starts the item type, not just the bare
+    /// single-character token.
     private mutating func sampleWhetherToCloseEmptyArray(
         items: GenerationSchema.Node
     ) async throws -> Bool {
-        let closeToken = try Self.singleToken(for: "]", backend: backend)
+        let closeTokens = tokensMatchingTrimmed("]")
         var allowed = try itemStartTokens(for: items)
-        allowed.insert(closeToken)
-        guard !allowed.isEmpty else {
+        allowed.formUnion(closeTokens)
+        guard !allowed.isEmpty, !closeTokens.isEmpty else {
             return false
         }
         let token = try await backend.sample(from: allowed)
-        return token == closeToken
+        return closeTokens.contains(token)
+    }
+
+    /// Tokens whose text equals `text` after trimming surrounding whitespace.
+    private func tokensMatchingTrimmed(_ text: String) -> Set<Int> {
+        var tokens = Set<Int>()
+        for token in 0 ..< backend.vocabSize {
+            if backend.isSpecialToken(token) { continue }
+            guard let tokenText = backend.tokenText(token) else { continue }
+            if tokenText.trimmingCharacters(in: .whitespacesAndNewlines) == text {
+                tokens.insert(token)
+            }
+        }
+        return tokens
+    }
+
+    /// Tokens whose text, after leading whitespace, is a non-empty prefix of one of `literals`.
+    private func tokensPrefixing(anyOf literals: [String]) -> Set<Int> {
+        var tokens = Set<Int>()
+        for token in 0 ..< backend.vocabSize {
+            if backend.isSpecialToken(token) { continue }
+            guard let text = backend.tokenText(token) else { continue }
+            let trimmed = text.drop(while: { $0.isWhitespace })
+            guard !trimmed.isEmpty else { continue }
+            if literals.contains(where: { $0.hasPrefix(trimmed) }) {
+                tokens.insert(token)
+            }
+        }
+        return tokens
+    }
+
+    /// Tokens whose first non-whitespace character is `prefix`.
+    private func tokensStarting(with prefix: Character) -> Set<Int> {
+        var tokens = Set<Int>()
+        for token in 0 ..< backend.vocabSize {
+            if backend.isSpecialToken(token) { continue }
+            guard let text = backend.tokenText(token) else { continue }
+            guard let first = text.drop(while: { $0.isWhitespace }).first else { continue }
+            if first == prefix {
+                tokens.insert(token)
+            }
+        }
+        return tokens
     }
 
     /// Tokens that can begin a JSON value for `node` (empty-array probe).
     private func itemStartTokens(for node: GenerationSchema.Node) throws -> Set<Int> {
         switch node {
         case .string:
-            return [quoteToken]
+            return tokensStarting(with: "\"")
         case .object:
-            return [try Self.singleToken(for: "{", backend: backend)]
+            return tokensStarting(with: "{")
         case .array:
-            return [try Self.singleToken(for: "[", backend: backend)]
+            return tokensStarting(with: "[")
         case .boolean:
-            var tokens = Set<Int>()
-            for literal in ["true", "false"] {
-                if let first = try backend.tokenize(literal).first {
-                    tokens.insert(first)
-                }
-            }
-            return tokens
+            return tokensPrefixing(anyOf: ["true", "false"])
         case .number(let numberNode):
             let numeric =
                 numberNode.integerOnly
@@ -621,9 +662,10 @@ struct ConstrainedJSONGenerator<Backend: TokenBackend> {
             // Only tokens that can start a number (digit or minus — not a bare `.`).
             return Set(
                 numeric.filter { token in
-                    guard let text = backend.tokenText(token), !text.isEmpty else { return false }
-                    let first = text.first
-                    return first?.isNumber == true || first == "-"
+                    guard let text = backend.tokenText(token),
+                        let first = text.drop(while: { $0.isWhitespace }).first
+                    else { return false }
+                    return first.isNumber || first == "-"
                 }
             )
         case .ref(let typeName):

Tests/AnyLanguageModelTests/StructuredGenerationTests.swift

@@ -588,6 +588,47 @@ struct StructuredGenerationTests {
         #expect(result == "[]")
     }
 
+    @Test func emptyArrayProbeAdmitsMergedTokens() async throws {
+        var maps = baseTokenMaps()
+        // Byte-pair vocabularies carry the item start and the close on merged tokens.
+        let spacedBracket = 60
+        let quoteA = 61
+        maps.tokenToText[spacedBracket] = " ]"
+        maps.tokenToText[quoteA] = "\"a"
+        let arrayNode = GenerationSchema.ArrayNode(
+            description: nil,
+            items: .string(.init(enumChoices: ["a"])),
+            minItems: nil,
+            maxItems: 1
+        )
+        let schema = GenerationSchema.primitive([String].self, node: .array(arrayNode))
+        let eosToken = 50
+
+        // The probe must offer the whitespace-prefixed close, and choosing it closes the array.
+        let closing = MockTokenBackend(
+            tokenToText: maps.tokenToText,
+            textToTokens: maps.textToTokens,
+            eosToken: eosToken,
+            endTokens: [eosToken],
+            maximumTokens: 64,
+            samplingQueue: [spacedBracket]
+        )
+        var closingGenerator = try ConstrainedJSONGenerator(backend: closing, schema: schema)
+        #expect(try await closingGenerator.generate() == "[]")
+
+        // The probe must offer the merged item start, and choosing it fills the array.
+        let filling = MockTokenBackend(
+            tokenToText: maps.tokenToText,
+            textToTokens: maps.textToTokens,
+            eosToken: eosToken,
+            endTokens: [eosToken],
+            maximumTokens: 64,
+            samplingQueue: [quoteA]
+        )
+        var fillingGenerator = try ConstrainedJSONGenerator(backend: filling, schema: schema)
+        #expect(try await fillingGenerator.generate() == "[\"a\"]")
+    }
+
     @Test func arrayTruncatesUnderBudgetPressure() async throws {
         let maps = baseTokenMaps()
         let arrayNode = GenerationSchema.ArrayNode(