From 46a4613e38af8cdf95ff2e27077f21dc78c2d807 Mon Sep 17 00:00:00 2001 From: Andrey Sapunov Date: Mon, 10 Aug 2026 13:49:13 +0200 Subject: [PATCH 1/3] fix: model-driven optional JSON object keys MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ConstrainedJSONGenerator pre-filtered optional properties with a hash of the field name XOR the token budget: let hash = key.utf8.reduce(0) { ($0 &* 31) &+ Int($1) } let combined = hash ^ backend.totalTokenBudget return combined % 2 == 0 Because 31 is odd and the budget is typically even, this reduces to the parity of the field name's bytes — identical for every document, decided before the model is consulted. Optional inclusion must not be a function of the key string. Keys are now chosen by masked sampling under the JSON grammar. After each property the mask permits any not-yet-emitted key, and permits `}` once every required property has been emitted. A last-resort budget floor still stops offering further optionals when remaining tokens are nearly exhausted and all required keys are already present. Visible behaviour change: for schemas with *no* required properties, `{}` is now reachable. That is schema-valid JSON and is what the model asked for when it samples the closing brace immediately. --- .../Shared/StructuredGeneration.swift | 108 ++++++++++-- .../StructuredGenerationTests.swift | 162 ++++++++++++++++++ 2 files changed, 252 insertions(+), 18 deletions(-) diff --git a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift index aae832e2..c626066d 100644 --- a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift +++ b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift @@ -20,10 +20,16 @@ protocol TokenBackend { var totalTokenBudget: Int { get } } -/// Heuristics for deciding when to include optional properties in generated output. -private enum OptionalPropertyBudget { - /// Minimum absolute number of tokens that should remain before we consider - /// adding optional properties. +/// Last-resort token-budget floor for optional structure (object keys / array items). +/// +/// Optional *selection* is model-driven (see ``ConstrainedJSONGenerator``). +/// This floor only kicks in when generation is about to run out of tokens: once the +/// schema-valid minimum has been emitted (all required keys, or `minItems` elements), +/// the generator stops offering further optionals and closes rather than risking a +/// hard budget failure mid-value. +private enum OptionalStructureBudget { + /// Minimum absolute number of tokens that should remain before offering more + /// optional properties or array elements. static let minimumRemainingTokens = 8 /// Require at least this fraction of the total budget (divisor form). @@ -425,17 +431,66 @@ struct ConstrainedJSONGenerator { } private mutating func generateObject(_ node: GenerationSchema.ObjectNode) async throws -> String { - let keys = node.properties.keys.sorted() - let includedKeys = keys.filter { shouldIncludeOptionalProperty($0, required: node.required) } + // Object *key set* is model-driven under the JSON grammar. The previous + // implementation pre-filtered optional properties with a hash of the field + // name XOR the token budget, so each optional was always-on or always-off + // for a given budget — decided before the model was consulted. + // + // After each property the mask permits any not-yet-emitted key, and permits + // `}` once every required property has been emitted. For schemas with no + // required properties that makes `{}` reachable (schema-valid, and what the + // model asked for) — a visible behaviour change for such schemas. + var remainingKeys = Set(node.properties.keys) + let required = node.required var output = try await emit("{") + var emittedAnyProperty = false + + while !remainingKeys.isEmpty { + let missingRequired = required.intersection(remainingKeys) + let canClose = missingRequired.isEmpty + let budgetAllowsMoreOptionals = hasBudgetForOptionalStructure() + + let keysToOffer: [String] + if budgetAllowsMoreOptionals { + // Model chooses any not-yet-emitted property (required or optional). + keysToOffer = remainingKeys.sorted() + } else if canClose { + // Genuine last resort: stop offering optionals when the budget is nearly + // exhausted. Required properties are already present, so close cleanly. + break + } else { + // Still missing required keys — only those may be emitted under pressure. + keysToOffer = missingRequired.sorted() + } - for (index, key) in includedKeys.enumerated() { - output += try await emit("\"\(key)\":") - output += try await generateNode(node.properties[key] ?? .string(.init())) + var candidates: [String] = keysToOffer.map { key in + let prefix = emittedAnyProperty ? "," : "" + return "\(prefix)\"\(key)\":" + } + // Closing is legal once every required property has been emitted. The model + // may leave remaining optionals out; that is schema-valid JSON. + if canClose { + candidates.append("}") + } - if index < includedKeys.count - 1 { - output += try await emit(",") + guard !candidates.isEmpty else { break } + + let choice = try await generateChoice(candidates) + output += choice + + // Model elected to close; remaining keys are optional and intentionally omitted. + if choice == "}" { + return output } + + guard let key = propertyKey(fromPropertyStart: choice), + let valueNode = node.properties[key] + else { + throw ConstrainedGenerationError.tokenizationFailed + } + remainingKeys.remove(key) + output += try await generateNode(valueNode) + emittedAnyProperty = true } output += try await emit("}") @@ -516,13 +571,30 @@ struct ConstrainedJSONGenerator { return output } - private func shouldIncludeOptionalProperty(_ key: String, required: Set) -> Bool { - if required.contains(key) { return true } - let minimumBudget = OptionalPropertyBudget.minimumBudget(totalTokenBudget: backend.totalTokenBudget) - guard backend.remainingTokens > minimumBudget else { return false } - let hash = key.utf8.reduce(0) { ($0 &* 31) &+ Int($1) } - let combined = hash ^ backend.totalTokenBudget - return combined % 2 == 0 + /// Whether enough budget remains to *offer* more optional structure to the model + /// (object properties or array elements). + /// + /// This is a last-resort guard only. It does not pick which optionals appear or how + /// long an array is — those decisions are made by constrained sampling. + private func hasBudgetForOptionalStructure() -> Bool { + let minimumBudget = OptionalStructureBudget.minimumBudget( + totalTokenBudget: backend.totalTokenBudget + ) + return backend.remainingTokens > minimumBudget + } + + /// Parses a property-start fragment produced for object key selection. + /// + /// Expected shapes: `"key":` (first property) or `,"key":` (subsequent). + private func propertyKey(fromPropertyStart choice: String) -> String? { + var fragment = choice + if fragment.first == "," { + fragment.removeFirst() + } + guard fragment.first == "\"", fragment.hasSuffix("\":") else { return nil } + fragment.removeFirst() + fragment.removeLast(2) + return fragment } private func deterministicChoice(from candidates: [String]) -> String { diff --git a/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift b/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift index 3b5f0b1b..86bcd1f7 100644 --- a/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift +++ b/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift @@ -450,4 +450,166 @@ struct StructuredGenerationTests { let result = try await generator.generate() #expect(result == "[\"a\",\"a\",\"a\"]") } + + // MARK: - Model-driven optional object properties + + private func objectTokenMaps() -> ( + tokenToText: [Int: String], + textToTokens: [String: [Int]] + ) { + // Structural + single-letter keys so `"x":` / `,"y":` tokenize without collisions. + var maps = baseTokenMaps() + let quote = 0 + let comma = 1 + let colon = 4 + let x = 10 + let y = 11 + let z = 12 + maps.textToTokens["\"x\":"] = [quote, x, quote, colon] + maps.textToTokens["\"y\":"] = [quote, y, quote, colon] + maps.textToTokens["\"z\":"] = [quote, z, quote, colon] + maps.textToTokens[",\"x\":"] = [comma, quote, x, quote, colon] + maps.textToTokens[",\"y\":"] = [comma, quote, y, quote, colon] + maps.textToTokens[",\"z\":"] = [comma, quote, z, quote, colon] + return maps + } + + private func allOptionalObjectSchema() -> GenerationSchema { + let stringNode = GenerationSchema.Node.string(.init(enumChoices: ["a"])) + let objectNode = GenerationSchema.ObjectNode( + description: nil, + properties: [ + "x": stringNode, + "y": stringNode, + "z": stringNode, + ], + required: [] + ) + // Type argument is unused; only the node shapes generation. + return GenerationSchema.primitive(String.self, node: .object(objectNode)) + } + + @Test func optionalObjectKeysAreChosenBySamplingNotNameHash() async throws { + let maps = objectTokenMaps() + let schema = allOptionalObjectSchema() + let eosToken = 50 + let quote = 0 + let y = 11 + let aToken = 8 + let colon = 4 + let rightBrace = 2 + + // Open with "y": (not lexicographically first), value "a", then close — leave x/z out. + // generateChoice samples every token of the chosen property-start and the enum value. + let backend = MockTokenBackend( + tokenToText: maps.tokenToText, + textToTokens: maps.textToTokens, + eosToken: eosToken, + endTokens: [eosToken], + maximumTokens: 64, + samplingQueue: [ + quote, y, quote, colon, // "y": + aToken, // enum value "a" + rightBrace, // close + ] + ) + + var generator = try ConstrainedJSONGenerator(backend: backend, schema: schema) + let result = try await generator.generate() + #expect(result == #"{"y":"a"}"#) + #expect(!result.contains("\"x\"")) + #expect(!result.contains("\"z\"")) + } + + @Test func optionalObjectKeysVaryWithSamplingQueue() async throws { + let maps = objectTokenMaps() + let schema = allOptionalObjectSchema() + let eosToken = 50 + let quote = 0 + let x = 10 + let z = 12 + let aToken = 8 + let comma = 1 + let colon = 4 + let rightBrace = 2 + + // Emit x then z (skip y) — different key set than the previous test. + let backend = MockTokenBackend( + tokenToText: maps.tokenToText, + textToTokens: maps.textToTokens, + eosToken: eosToken, + endTokens: [eosToken], + maximumTokens: 64, + samplingQueue: [ + quote, x, quote, colon, // "x": + aToken, // "a" + comma, quote, z, quote, colon, // ,"z": + aToken, // "a" + rightBrace, + ] + ) + + var generator = try ConstrainedJSONGenerator(backend: backend, schema: schema) + let result = try await generator.generate() + #expect(result == #"{"x":"a","z":"a"}"#) + #expect(!result.contains("\"y\"")) + } + + @Test func requiredObjectKeysMustBeEmittedBeforeClose() async throws { + let maps = objectTokenMaps() + let stringNode = GenerationSchema.Node.string(.init(enumChoices: ["a"])) + let objectNode = GenerationSchema.ObjectNode( + description: nil, + properties: [ + "x": stringNode, + "y": stringNode, + ], + required: ["x"] + ) + let schema = GenerationSchema.primitive(String.self, node: .object(objectNode)) + let eosToken = 50 + let quote = 0 + let x = 10 + let aToken = 8 + let colon = 4 + let rightBrace = 2 + + // "}" is not among candidates until required "x" is emitted. + let backend = MockTokenBackend( + tokenToText: maps.tokenToText, + textToTokens: maps.textToTokens, + eosToken: eosToken, + endTokens: [eosToken], + maximumTokens: 64, + samplingQueue: [ + quote, x, quote, colon, + aToken, + rightBrace, + ] + ) + + var generator = try ConstrainedJSONGenerator(backend: backend, schema: schema) + let result = try await generator.generate() + #expect(result == #"{"x":"a"}"#) + } + + @Test func emptyObjectWhenModelClosesImmediately() async throws { + let maps = objectTokenMaps() + let schema = allOptionalObjectSchema() + let eosToken = 50 + let rightBrace = 2 + + let backend = MockTokenBackend( + tokenToText: maps.tokenToText, + textToTokens: maps.textToTokens, + eosToken: eosToken, + endTokens: [eosToken], + maximumTokens: 64, + samplingQueue: [rightBrace] + ) + + var generator = try ConstrainedJSONGenerator(backend: backend, schema: schema) + let result = try await generator.generate() + #expect(result == "{}") + } } From 7b2705624f944288d3b0d0e5b391b35aae190cbd Mon Sep 17 00:00:00 2001 From: Andrey Sapunov Date: Mon, 10 Aug 2026 13:50:34 +0200 Subject: [PATCH 2/3] fix: model-driven array length under the JSON grammar MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ConstrainedJSONGenerator derived array element counts from the token budget before generation: let budgetBasedCount = backend.totalTokenBudget / arrayDefaultCountDivisor ... let offset = rangeSize > 0 ? backend.totalTokenBudget % rangeSize : 0 count = minItems + offset The loop then emitted exactly that many elements, with no way to stop early. A three-item list was padded with invented elements; a ten-item list was truncated. Same defect shape as pre-filtering optional object keys by name hash: the length was a function of a constant budget, identical for every document, decided before the model was consulted. Array length is now chosen by masked sampling. After each element the mask permits both `,` and `]`, honouring schema `minItems` / `maxItems` when stated. Empty arrays (`minItems == 0` / unbounded) are reachable by sampling `]` immediately after `[`. A last-resort budget floor still force-closes once `minItems` is satisfied when remaining tokens are nearly exhausted. Drops the `totalTokenBudget % rangeSize` arithmetic entirely — an arbitrary count inside a valid range is still arbitrary. --- .../Shared/StructuredGeneration.swift | 145 +++++++++++--- .../StructuredGenerationTests.swift | 179 +++++++++++++++++- 2 files changed, 291 insertions(+), 33 deletions(-) diff --git a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift index c626066d..f26b0ea5 100644 --- a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift +++ b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift @@ -75,10 +75,6 @@ struct ConstrainedJSONGenerator { private static var maxIntegerTokenLimit: Int { 20 } private static var maxDecimalTokenLimit: Int { 32 } - /// Heuristics for default array sizes when no bounds are specified. - private static var arrayDefaultCountDivisor: Int { 32 } - private static var arrayDefaultCountMax: Int { 16 } - private var backend: Backend private let schema: GenerationSchema private var emittedText = "" @@ -498,42 +494,131 @@ struct ConstrainedJSONGenerator { } private mutating func generateArray(_ node: GenerationSchema.ArrayNode) async throws -> String { - // Derive a default item count from the total token budget when the schema - // does not specify explicit minItems/maxItems. We use a small fraction of the - // budget and clamp it to a reasonable range to avoid overlong arrays. - let budgetBasedCount = backend.totalTokenBudget / Self.arrayDefaultCountDivisor - let defaultCount = max(1, min(Self.arrayDefaultCountMax, budgetBasedCount)) - let count: Int - - if let minItems = node.minItems, let maxItems = node.maxItems { - if minItems > maxItems { - throw ConstrainedGenerationError.invalidArrayBounds( - "Minimum items \(minItems) exceeds maximum \(maxItems)" - ) - } - let rangeSize = maxItems - minItems + 1 - let offset = rangeSize > 0 ? backend.totalTokenBudget % rangeSize : 0 - count = minItems + offset - } else if let minItems = node.minItems { - count = minItems - } else if let maxItems = node.maxItems { - count = maxItems - } else { - count = defaultCount + // Array *length* is model-driven under the JSON grammar — same family of fix as + // model-driven optional object keys. After each element the mask permits both + // continuing (`,`) and closing (`]`), subject to schema `minItems` / `maxItems`. + // A budget-derived fixed count (or `totalTokenBudget % rangeSize`) would force the + // same length for every document and invent filler or truncate real items. + let minItems = max(0, node.minItems ?? 0) + let maxItems = node.maxItems + + if let maxItems, minItems > maxItems { + throw ConstrainedGenerationError.invalidArrayBounds( + "Minimum items \(minItems) exceeds maximum \(maxItems)" + ) } + var output = try await emit("[") + var count = 0 - for index in 0 ..< count { - output += try await generateNode(node.items) - if index < count - 1 { - output += try await emit(",") + while true { + if let maxItems, count >= maxItems { + break } + + let canClose = count >= minItems + let budgetAllowsMore = hasBudgetForOptionalStructure() + + if canClose && !budgetAllowsMore { + // Genuine last resort: close once minItems is satisfied rather than + // failing mid-element under a hard budget floor. + break + } + + if count > 0 { + if canClose { + // Model chooses continue vs close. + let choice = try await generateChoice([",", "]"]) + if choice == "]" { + output += choice + return output + } + output += choice + } else { + // Still below minItems — must emit another element. + output += try await emit(",") + } + } else if canClose { + // Empty array is legal (`minItems == 0`). Probe whether the model wants + // `]` or a first element. Sampling is non-committing for non-`]` tokens + // (see ``sampleWhetherToCloseEmptyArray``); the element is then generated + // from the same decode state. + if try await sampleWhetherToCloseEmptyArray(items: node.items) { + output += try await emit("]") + return output + } + } + + output += try await generateNode(node.items) + count += 1 } output += try await emit("]") return output } + /// 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. + private mutating func sampleWhetherToCloseEmptyArray( + items: GenerationSchema.Node + ) async throws -> Bool { + let closeToken = try Self.singleToken(for: "]", backend: backend) + var allowed = try itemStartTokens(for: items) + allowed.insert(closeToken) + guard !allowed.isEmpty else { + return false + } + let token = try await backend.sample(from: allowed) + return token == closeToken + } + + /// Tokens that can begin a JSON value for `node` (empty-array probe). + private func itemStartTokens(for node: GenerationSchema.Node) throws -> Set { + switch node { + case .string: + return [quoteToken] + case .object: + return [try Self.singleToken(for: "{", backend: backend)] + case .array: + return [try Self.singleToken(for: "[", backend: backend)] + case .boolean: + var tokens = Set() + for literal in ["true", "false"] { + if let first = try backend.tokenize(literal).first { + tokens.insert(first) + } + } + return tokens + case .number(let numberNode): + let numeric = + numberNode.integerOnly + ? integerTerminators.subtracting(basicTerminators) + : doubleTerminators.subtracting(basicTerminators) + // 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 == "-" + } + ) + case .ref(let typeName): + guard let referenced = schema.defs[typeName] else { + throw ConstrainedGenerationError.missingReference(typeName) + } + return try itemStartTokens(for: referenced) + case .anyOf(let variants): + var tokens = Set() + for variant in variants { + tokens.formUnion(try itemStartTokens(for: variant)) + } + return tokens + } + } + private mutating func generateString(_ node: GenerationSchema.StringNode) async throws -> String { var output = try await emit("\"") let content: String diff --git a/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift b/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift index 86bcd1f7..565961a7 100644 --- a/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift +++ b/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift @@ -428,7 +428,9 @@ struct StructuredGenerationTests { } } - @Test func arrayCountIsDeterministic() async throws { + // MARK: - Model-driven array length + + @Test func arrayLengthIsChosenBySamplingNotBudget() async throws { let maps = baseTokenMaps() let arrayNode = GenerationSchema.ArrayNode( description: nil, @@ -438,17 +440,188 @@ struct StructuredGenerationTests { ) let schema = GenerationSchema.primitive([String].self, node: .array(arrayNode)) let eosToken = 50 + let aToken = 8 + let rightBracket = 3 + + // minItems=1 forces first element; model then closes (length 1), not budget-derived 3. + // With maximumTokens 17 the old formula was minItems + (17 % 3) = 3. let backend = MockTokenBackend( tokenToText: maps.tokenToText, textToTokens: maps.textToTokens, eosToken: eosToken, endTokens: [eosToken], - maximumTokens: 17 + maximumTokens: 17, + samplingQueue: [ + aToken, // first "a" + rightBracket, // close after 1 (`,` would continue) + ] + ) + + var generator = try ConstrainedJSONGenerator(backend: backend, schema: schema) + let result = try await generator.generate() + #expect(result == "[\"a\"]") + } + + @Test func arrayLengthVariesWithSamplingQueue() async throws { + let maps = baseTokenMaps() + let arrayNode = GenerationSchema.ArrayNode( + description: nil, + items: .string(.init(enumChoices: ["a"])), + minItems: 1, + maxItems: 3 + ) + let schema = GenerationSchema.primitive([String].self, node: .array(arrayNode)) + let eosToken = 50 + let aToken = 8 + let comma = 1 + let rightBracket = 3 + + // Emit two elements then close — different length than the previous test. + let backend = MockTokenBackend( + tokenToText: maps.tokenToText, + textToTokens: maps.textToTokens, + eosToken: eosToken, + endTokens: [eosToken], + maximumTokens: 64, + samplingQueue: [ + aToken, // "a" + comma, // continue + aToken, // "a" + rightBracket, // close + ] + ) + + var generator = try ConstrainedJSONGenerator(backend: backend, schema: schema) + let result = try await generator.generate() + #expect(result == "[\"a\",\"a\"]") + } + + @Test func arrayRespectsMaxItems() async throws { + let maps = baseTokenMaps() + let arrayNode = GenerationSchema.ArrayNode( + description: nil, + items: .string(.init(enumChoices: ["a"])), + minItems: 1, + maxItems: 2 + ) + let schema = GenerationSchema.primitive([String].self, node: .array(arrayNode)) + let eosToken = 50 + let aToken = 8 + let comma = 1 + + // Model always continues; generator must still stop at maxItems=2. + let backend = MockTokenBackend( + tokenToText: maps.tokenToText, + textToTokens: maps.textToTokens, + eosToken: eosToken, + endTokens: [eosToken], + maximumTokens: 64, + samplingQueue: [ + aToken, + comma, + aToken, + // further commas would be illegal once max is reached — close is forced + ] + ) + + var generator = try ConstrainedJSONGenerator(backend: backend, schema: schema) + let result = try await generator.generate() + #expect(result == "[\"a\",\"a\"]") + } + + @Test func arrayRespectsMinItemsBeforeClose() async throws { + let maps = baseTokenMaps() + let arrayNode = GenerationSchema.ArrayNode( + description: nil, + items: .string(.init(enumChoices: ["a"])), + minItems: 2, + maxItems: 4 + ) + let schema = GenerationSchema.primitive([String].self, node: .array(arrayNode)) + let eosToken = 50 + let aToken = 8 + let rightBracket = 3 + + // After first element, `]` is not offered — only forced `,` + second element, then close. + let backend = MockTokenBackend( + tokenToText: maps.tokenToText, + textToTokens: maps.textToTokens, + eosToken: eosToken, + endTokens: [eosToken], + maximumTokens: 64, + samplingQueue: [ + aToken, // first + // no close offered here — comma is emitted forcibly + aToken, // second (satisfies minItems) + rightBracket, // model closes + ] + ) + + var generator = try ConstrainedJSONGenerator(backend: backend, schema: schema) + let result = try await generator.generate() + #expect(result == "[\"a\",\"a\"]") + } + + @Test func emptyArrayWhenModelClosesImmediately() async throws { + let maps = baseTokenMaps() + let arrayNode = GenerationSchema.ArrayNode( + description: nil, + items: .string(.init(enumChoices: ["a"])), + minItems: nil, + maxItems: nil + ) + let schema = GenerationSchema.primitive([String].self, node: .array(arrayNode)) + let eosToken = 50 + let rightBracket = 3 + + let backend = MockTokenBackend( + tokenToText: maps.tokenToText, + textToTokens: maps.textToTokens, + eosToken: eosToken, + endTokens: [eosToken], + maximumTokens: 64, + samplingQueue: [rightBracket] + ) + + var generator = try ConstrainedJSONGenerator(backend: backend, schema: schema) + let result = try await generator.generate() + #expect(result == "[]") + } + + @Test func arrayTruncatesUnderBudgetPressure() async throws { + let maps = baseTokenMaps() + let arrayNode = GenerationSchema.ArrayNode( + description: nil, + items: .string(.init(enumChoices: ["a"])), + minItems: 1, + maxItems: 8 + ) + let schema = GenerationSchema.primitive([String].self, node: .array(arrayNode)) + let eosToken = 50 + let aToken = 8 + let comma = 1 + + // Mock maps encode `]` / `"` / `a` but not `[`, so the opening bracket is free. + // First element costs 3 tokens (`"`, `a`, `"`). Floor is max(8, budget/10)=8. + // With budget 11, remaining after the first element is 8 → not strictly greater + // than the floor → force-close. Sampling queue would happily continue with `,`. + let backend = MockTokenBackend( + tokenToText: maps.tokenToText, + textToTokens: maps.textToTokens, + eosToken: eosToken, + endTokens: [eosToken], + maximumTokens: 11, + samplingQueue: [ + aToken, + comma, // would continue if offered — must not be consumed if we truncate + aToken, + ] ) var generator = try ConstrainedJSONGenerator(backend: backend, schema: schema) let result = try await generator.generate() - #expect(result == "[\"a\",\"a\",\"a\"]") + // Force-close after satisfying minItems under budget pressure. + #expect(result == "[\"a\"]") } // MARK: - Model-driven optional object properties From a3097b6413eaa4e9e0bdef0c38b45c43a6b0361c Mon Sep 17 00:00:00 2001 From: Andrey Sapunov Date: Mon, 10 Aug 2026 13:51:27 +0200 Subject: [PATCH 3/3] fix: allow standalone '.' and '-' in constrained number tokens buildValidDecimalTokens required every candidate token to contain a digit: if text.allSatisfy({ $0.isNumber || $0 == "-" || $0 == "." }), text.contains(where: { $0.isNumber }) Tokenizers that encode `.` and `-` as standalone tokens (Qwen2.5 does; `473.00` is `4` `7` `3` `.` `0` `0`) therefore could not produce a decimal point at all. Observed effect: the model emits the integer part, cannot place the point, keeps emitting digits until the token cap, and the result re-serialises as something like `4.73e+31`. Include standalone `.` / `-` in the decimal mask (and standalone `-` in the integer mask). Restrict accepted digits to ASCII `0`...`9` so fullwidth / superscript forms that `Character.isNumber` admits cannot enter JSON numbers. --- .../Shared/StructuredGeneration.swift | 33 +++++-- .../StructuredGenerationTests.swift | 86 +++++++++++++++++++ 2 files changed, 113 insertions(+), 6 deletions(-) diff --git a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift index f26b0ea5..e02d88ab 100644 --- a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift +++ b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift @@ -195,28 +195,49 @@ struct ConstrainedJSONGenerator { return samples.filter { $0 >= 0 && $0 < vocabSize }.sorted() } + /// ASCII digit `0`...`9` only — JSON numbers must not accept fullwidth / superscript + /// forms that `Character.isNumber` would otherwise admit. + private static func isASCIIDigit(_ character: Character) -> Bool { + character >= "0" && character <= "9" + } + + /// Tokens that may appear inside a JSON integer: ASCII digits and standalone `-`. + /// + /// Standalone `-` is required because BPE tokenizers (Qwen2.5, etc.) encode `-1` as + /// two tokens. Requiring every token to contain a digit excluded `-` and made negatives + /// unrepresentable except via rare multi-character tokens. private static func buildValidIntegerTokens(backend: Backend) -> Set { var allowed = Set() for token in 0 ..< backend.vocabSize { if backend.isSpecialToken(token) { continue } guard let text = backend.tokenText(token), !text.isEmpty else { continue } - if text.allSatisfy({ $0.isNumber || $0 == "-" }), - text.contains(where: { $0.isNumber }) - { + let onlyIntegerChars = text.allSatisfy { Self.isASCIIDigit($0) || $0 == "-" } + let hasDigit = text.contains { Self.isASCIIDigit($0) } + let isStandaloneMinus = text == "-" + if onlyIntegerChars && (hasDigit || isStandaloneMinus) { allowed.insert(token) } } return allowed } + /// Tokens that may appear inside a JSON number: ASCII digits, `-`, and `.`. + /// + /// **Critical:** standalone `.` and `-` must be included. Qwen2.5 encodes `473.00` as + /// `4` `7` `3` `.` `0` `0`. The previous filter required every token to contain a digit, + /// which dropped `.` and forced the model to pad zeros until `maxDecimalTokenLimit` + /// (pathological `e+31` values after Double re-serialization). private static func buildValidDecimalTokens(backend: Backend) -> Set { var allowed = Set() for token in 0 ..< backend.vocabSize { if backend.isSpecialToken(token) { continue } guard let text = backend.tokenText(token), !text.isEmpty else { continue } - if text.allSatisfy({ $0.isNumber || $0 == "-" || $0 == "." }), - text.contains(where: { $0.isNumber }) - { + let onlyNumberChars = text.allSatisfy { + Self.isASCIIDigit($0) || $0 == "-" || $0 == "." + } + let hasDigit = text.contains { Self.isASCIIDigit($0) } + let isStandaloneSignOrDot = text == "-" || text == "." + if onlyNumberChars && (hasDigit || isStandaloneSignOrDot) { allowed.insert(token) } } diff --git a/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift b/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift index 565961a7..eff70b83 100644 --- a/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift +++ b/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift @@ -785,4 +785,90 @@ struct StructuredGenerationTests { let result = try await generator.generate() #expect(result == "{}") } + + // MARK: - Decimal / number token mask + + private func numberTokenMaps() -> ( + tokenToText: [Int: String], + textToTokens: [String: [Int]] + ) { + var maps = baseTokenMaps() + let dotToken = 20 + maps.tokenToText[dotToken] = "." + maps.textToTokens["."] = [dotToken] + return maps + } + + @Test func decimalNumberEmitsStandaloneDot() async throws { + // Qwen2.5-style tokenization of 473.00 is 4 7 3 . 0 0. Standalone `.` must be + // in the decimal mask; otherwise the model cannot place the point and pads digits + // until the token cap (re-serialized as e+31 after Double conversion). + var maps = numberTokenMaps() + maps.tokenToText[30] = "4" + maps.tokenToText[31] = "7" + maps.tokenToText[32] = "3" + maps.textToTokens["4"] = [30] + maps.textToTokens["7"] = [31] + maps.textToTokens["3"] = [32] + let numberNode = GenerationSchema.NumberNode( + description: nil, + minimum: nil, + maximum: nil, + integerOnly: false + ) + let schema = GenerationSchema.primitive(Double.self, node: .number(numberNode)) + let eosToken = 50 + let fourToken = 30 + let sevenToken = 31 + let threeToken = 32 + let dotToken = 20 + let zeroToken = 5 + let rightBrace = 2 + + let backend = MockTokenBackend( + tokenToText: maps.tokenToText, + textToTokens: maps.textToTokens, + eosToken: eosToken, + endTokens: [eosToken], + maximumTokens: 64, + samplingQueue: [ + fourToken, sevenToken, threeToken, // 473 + dotToken, zeroToken, zeroToken, // .00 + rightBrace, // terminate + ] + ) + + var generator = try ConstrainedJSONGenerator(backend: backend, schema: schema) + let result = try await generator.generate() + #expect(result == "473.00") + } + + @Test func standaloneMinusIsAllowedInIntegerMask() async throws { + // Standalone `-` then digit, as BPE tokenizers encode negatives. + let maps = baseTokenMaps() + let numberNode = GenerationSchema.NumberNode( + description: nil, + minimum: -10, + maximum: 0, + integerOnly: true + ) + let schema = GenerationSchema.primitive(Int.self, node: .number(numberNode)) + let eosToken = 50 + let minusToken = 13 + let oneToken = 6 + let rightBrace = 2 + + let backend = MockTokenBackend( + tokenToText: maps.tokenToText, + textToTokens: maps.textToTokens, + eosToken: eosToken, + endTokens: [eosToken], + maximumTokens: 8, + samplingQueue: [minusToken, oneToken, rightBrace] + ) + + var generator = try ConstrainedJSONGenerator(backend: backend, schema: schema) + let result = try await generator.generate() + #expect(result == "-1") + } }