diff --git a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift index aae832e2..e02d88ab 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). @@ -69,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 = "" @@ -193,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) } } @@ -425,17 +448,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() + } + + 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("}") + } - for (index, key) in includedKeys.enumerated() { - output += try await emit("\"\(key)\":") - output += try await generateNode(node.properties[key] ?? .string(.init())) + 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 + } - if index < includedKeys.count - 1 { - output += try await emit(",") + 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("}") @@ -443,42 +515,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 @@ -516,13 +677,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..eff70b83 100644 --- a/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift +++ b/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift @@ -428,7 +428,41 @@ 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, + 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 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, + 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, @@ -438,16 +472,403 @@ struct StructuredGenerationTests { ) 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() + // Force-close after satisfying minItems under budget pressure. + #expect(result == "[\"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: 17 + maximumTokens: 64, + samplingQueue: [rightBrace] ) var generator = try ConstrainedJSONGenerator(backend: backend, schema: schema) let result = try await generator.generate() - #expect(result == "[\"a\",\"a\",\"a\"]") + #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") } }