From 3c7a0bdffcd3edcafa4bbe673b4b5a53770cbaff Mon Sep 17 00:00:00 2001 From: Jaleel Akbashev Date: Tue, 6 Oct 2026 10:50:07 +0200 Subject: [PATCH 1/9] Support payloads in Generable enums --- .../AnyLanguageModel/GenerationSchema.swift | 49 +++ .../GenerableMacro.swift | 320 +++--------------- .../GenerableEnumPayloadTests.swift | 96 ++++++ .../SystemLanguageModelTests.swift | 26 ++ 4 files changed, 223 insertions(+), 268 deletions(-) create mode 100644 Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift diff --git a/Sources/AnyLanguageModel/GenerationSchema.swift b/Sources/AnyLanguageModel/GenerationSchema.swift index 9df4978c..e9e6f3c1 100644 --- a/Sources/AnyLanguageModel/GenerationSchema.swift +++ b/Sources/AnyLanguageModel/GenerationSchema.swift @@ -397,6 +397,55 @@ public struct GenerationSchema: Equatable, Codable, CustomDebugStringConvertible self.defs = [:] } + /// Creates a schema for an enumeration whose cases carry values, the way Foundation Models + /// describes one: one object for each case, with a `type` property that names the case and a + /// property for each of its values. + /// + /// - Parameters: + /// - type: The type this schema represents. + /// - description: A natural language description of this schema, given to each case. + /// - cases: Each case's name, with a property for each of its values. + public init( + type: any Generable.Type, + description: String? = nil, + cases: [(name: String, properties: [GenerationSchema.Property])] + ) { + let typeName = String(reflecting: type) + var members: [Node] = [] + var allDefs: [String: Node] = [:] + + for (name, properties) in cases { + var props: [String: Node] = ["type": .string(StringNode(enumChoices: [name]))] + var required: Set = ["type"] + for property in properties { + props[property.name] = property.node + if !property.isOptional { + required.insert(property.name) + } + for (defName, defNode) in property.deps { + if let existing = allDefs[defName], existing != defNode { + fatalError("Duplicate type '\(defName)' with different structure") + } + allDefs[defName] = defNode + } + } + let caseName = "\(typeName).Discriminated\(name.prefix(1).uppercased() + name.dropFirst())" + allDefs[caseName] = .object( + ObjectNode( + description: description, + properties: props, + required: required, + propertyOrder: ["type"] + properties.map(\.name) + ) + ) + members.append(.ref(caseName)) + } + + allDefs[typeName] = .anyOf(members) + self.root = .ref(typeName) + self.defs = allDefs + } + /// Creates a schema as the union of several other types. /// /// - Parameters: diff --git a/Sources/AnyLanguageModelMacros/GenerableMacro.swift b/Sources/AnyLanguageModelMacros/GenerableMacro.swift index 1e810ad8..d9435fbc 100644 --- a/Sources/AnyLanguageModelMacros/GenerableMacro.swift +++ b/Sources/AnyLanguageModelMacros/GenerableMacro.swift @@ -1070,68 +1070,31 @@ public struct GenerableMacro: MemberMacro, ExtensionMacro { let hasAnyAssociatedValues = cases.contains { $0.hasAssociatedValues } if hasAnyAssociatedValues { + // Each case is an object whose `type` names it, with a property for each value, + // as Foundation Models generates it. let switchCases = cases.map { enumCase in - if enumCase.associatedValues.isEmpty { - return """ - case "\(enumCase.name)": - self = .\(enumCase.name) - """ - } else if enumCase.isSingleUnlabeledValue { - let valueType = enumCase.associatedValues[0].type - return generateSingleValueCase(caseName: enumCase.name, valueType: valueType) - } else { - return generateMultipleValueCase( - caseName: enumCase.name, - associatedValues: enumCase.associatedValues - ) + let arguments = enumCase.payload.map { value in + let read = "try generatedContent.value(\(value.type).self, forProperty: \"\(value.property)\")" + return value.label.map { "\($0): \(read)" } ?? read } + let construction = + arguments.isEmpty ? ".\(enumCase.name)" : ".\(enumCase.name)(\(arguments.joined(separator: ", ")))" + return """ + case "\(enumCase.name)": + self = \(construction) + """ }.joined(separator: "\n ") return DeclSyntax( stringLiteral: """ nonisolated public init(_ generatedContent: GeneratedContent) throws { - do { - guard case .structure(let properties, _) = generatedContent.kind else { - throw DecodingError.typeMismatch( - \(enumName).self, - DecodingError.Context(codingPath: [], debugDescription: "Expected structure for enum \(enumName)") - ) - } - - guard case .string(let caseValue) = properties["case"]?.kind else { - struct Key: CodingKey { - var stringValue: String - var intValue: Int? { nil } - init(stringValue: String) { self.stringValue = stringValue } - init?(intValue: Int) { nil } - } - throw DecodingError.keyNotFound( - Key(stringValue: "case"), - DecodingError.Context(codingPath: [], debugDescription: "Missing 'case' property in enum data for \(enumName)") - ) - } - - let valueContent = properties["value"] - - switch caseValue { - \(switchCases) - default: - throw DecodingError.dataCorrupted( - DecodingError.Context(codingPath: [], debugDescription: "Invalid enum case '\\(caseValue)' for \(enumName). Valid cases: [\(cases.map { $0.name }.joined(separator: ", "))]") - ) - } - } catch { - guard case .string(let value) = generatedContent.kind else { - throw error - } - let trimmedValue = value.trimmingCharacters(in: .whitespacesAndNewlines) - switch trimmedValue { - \(cases.filter { !$0.hasAssociatedValues }.map { "case \"\($0.name)\": self = .\($0.name)" }.joined(separator: "\n ")) - default: - throw DecodingError.dataCorrupted( - DecodingError.Context(codingPath: [], debugDescription: "Invalid enum case '\\(trimmedValue)' for \(enumName). Valid cases: [\(cases.map { $0.name }.joined(separator: ", "))]") - ) - } + let type = try generatedContent.value(String.self, forProperty: "type") + switch type { + \(switchCases) + default: + throw DecodingError.dataCorrupted( + DecodingError.Context(codingPath: [], debugDescription: "Invalid enum case '\\(type)' for \(enumName). Valid cases: [\(cases.map { $0.name }.joined(separator: ", "))]") + ) } } """ @@ -1165,116 +1128,6 @@ public struct GenerableMacro: MemberMacro, ExtensionMacro { } } - private static func generateSingleValueCase(caseName: String, valueType: String) -> String { - switch valueType { - case "String": - return """ - case "\(caseName)": - if let valueContent = valueContent, - case .string(let stringValue) = valueContent.kind { - self = .\(caseName)(stringValue) - } else { - self = .\(caseName)("") - } - """ - case "Int": - return """ - case "\(caseName)": - if let valueContent = valueContent { - let intValue = try valueContent.value(Int.self) - self = .\(caseName)(intValue) - } else { - self = .\(caseName)(0) - } - """ - case "Double": - return """ - case "\(caseName)": - if let valueContent = valueContent { - let doubleValue = try valueContent.value(Double.self) - self = .\(caseName)(doubleValue) - } else { - self = .\(caseName)(0.0) - } - """ - case "Bool": - return """ - case "\(caseName)": - if let valueContent = valueContent { - let boolValue = try valueContent.value(Bool.self) - self = .\(caseName)(boolValue) - } else { - self = .\(caseName)(false) - } - """ - default: - return """ - case "\(caseName)": - if let valueContent = valueContent { - let associatedValue = try \(valueType)(valueContent) - self = .\(caseName)(associatedValue) - } else { - throw DecodingError.valueNotFound( - \(valueType).self, - DecodingError.Context(codingPath: [], debugDescription: "Missing value for enum case '\(caseName)' with associated type \(valueType)") - ) - } - """ - } - } - - private static func generateMultipleValueCase( - caseName: String, - associatedValues: [(label: String?, type: String)] - ) -> String { - let valueExtractions = associatedValues.enumerated().map { index, assocValue in - let label = assocValue.label ?? "param\(index)" - let type = assocValue.type - - switch type { - case "String": - return "let \(label) = try valueProperties[\"\(label)\"]?.value(String.self) ?? \"\"" - case "Int": - return "let \(label) = try valueProperties[\"\(label)\"]?.value(Int.self) ?? 0" - case "Double": - return "let \(label) = try valueProperties[\"\(label)\"]?.value(Double.self) ?? 0.0" - case "Bool": - return "let \(label) = try valueProperties[\"\(label)\"]?.value(Bool.self) ?? false" - default: - return - "let \(label) = try \(type)(valueProperties[\"\(label)\"] ?? GeneratedContent(\"{}\"))" - } - }.joined(separator: "\n ") - - let parameterList = associatedValues.enumerated().map { index, assocValue in - let label = assocValue.label ?? "param\(index)" - if assocValue.label != nil { - return "\(label): \(label)" - } else { - return label - } - }.joined(separator: ", ") - - return """ - case "\(caseName)": - if let valueContent = valueContent { - guard case .structure(let valueProperties, _) = valueContent.kind else { - throw DecodingError.typeMismatch( - [String: Any].self, - DecodingError.Context(codingPath: [], debugDescription: "Expected structure for enum case '\(caseName)' associated values") - ) - } - \(valueExtractions) - self = .\(caseName)(\(parameterList)) - } else { - throw DecodingError.valueNotFound( - [String: Any].self, - DecodingError.Context(codingPath: [], debugDescription: "Missing value data for enum case '\(caseName)' with associated values") - ) - } - """ - } - private static func generateEnumGeneratedContentProperty( enumName: String, description: String?, @@ -1284,28 +1137,25 @@ public struct GenerableMacro: MemberMacro, ExtensionMacro { if hasAnyAssociatedValues { let switchCases = cases.map { enumCase in - if enumCase.associatedValues.isEmpty { - return """ - case .\(enumCase.name): - return GeneratedContent(properties: [ - "case": GeneratedContent("\(enumCase.name)"), - "value": GeneratedContent("") - ]) - """ - } else if enumCase.isSingleUnlabeledValue { - return """ - case .\\(enumCase.name)(let value): - return GeneratedContent(properties: [ - "case": GeneratedContent("\\(enumCase.name)"), - "value": GeneratedContent("\\\\(value)") - ]) + let payload = enumCase.payload + let bindings = payload.indices.map { "let value\($0)" }.joined(separator: ", ") + let pattern = payload.isEmpty ? ".\(enumCase.name)" : ".\(enumCase.name)(\(bindings))" + // An optional value that's nil is left out, as for a structure's properties. + let assignments = payload.enumerated().map { index, value in + let assignment = """ + properties["\(value.property)"] = value\(index).generatedContent + keys.append("\(value.property)") """ - } else { - return generateMultipleValueSerialization( - caseName: enumCase.name, - associatedValues: enumCase.associatedValues - ) - } + return value.isOptional + ? "if let value\(index) {\n \(assignment)\n }" : assignment + }.joined(separator: "\n ") + return """ + case \(pattern): + var properties: [String: GeneratedContent] = ["type": GeneratedContent("\(enumCase.name)")] + var keys = ["type"] + \(assignments) + return GeneratedContent(kind: .structure(properties: properties, orderedKeys: keys)) + """ }.joined(separator: "\n ") return DeclSyntax( @@ -1334,61 +1184,6 @@ public struct GenerableMacro: MemberMacro, ExtensionMacro { } } - private static func generateSingleValueSerialization(caseName: String, valueType: String) - -> String - { - switch valueType { - case "String", "Int", "Double", "Bool": - return """ - case .\(caseName)(let value): - return GeneratedContent(properties: [ - "case": GeneratedContent("\(caseName)"), - "value": GeneratedContent("\\(value)") - ]) - """ - default: - return """ - case .\(caseName)(let value): - return GeneratedContent(properties: [ - "case": GeneratedContent("\(caseName)"), - "value": value.generatedContent - ]) - """ - } - } - - private static func generateMultipleValueSerialization( - caseName: String, - associatedValues: [(label: String?, type: String)] - ) -> String { - let parameterList = associatedValues.enumerated().map { index, assocValue in - let label = assocValue.label ?? "param\(index)" - return "let \(label)" - }.joined(separator: ", ") - - let propertyMappings = associatedValues.enumerated().map { index, assocValue in - let label = assocValue.label ?? "param\(index)" - let type = assocValue.type - - switch type { - case "String", "Int", "Double", "Bool": - return "\"\(label)\": GeneratedContent(\"\\(\(label))\")" - default: - return "\"\(label)\": \(label).generatedContent" - } - }.joined(separator: ",\n ") - - return """ - case .\(caseName)(\(parameterList)): - return GeneratedContent(properties: [ - "case": GeneratedContent("\(caseName)"), - "value": GeneratedContent(properties: [ - \(propertyMappings) - ]) - ]) - """ - } - private static func generateEnumFromGeneratedContentMethod(enumName: String) -> DeclSyntax { return DeclSyntax( stringLiteral: """ @@ -1407,23 +1202,12 @@ public struct GenerableMacro: MemberMacro, ExtensionMacro { let hasAnyAssociatedValues = cases.contains { $0.hasAssociatedValues } if hasAnyAssociatedValues { - - let caseProperty = """ - GenerationSchema.Property( - name: "case", - description: "Enum case identifier", - type: String.self, - guides: [] - ) - """ - let valueProperty = """ - GenerationSchema.Property( - name: "value", - description: "Associated value data", - type: String.self, - guides: [] - ) - """ + let caseSchemas = cases.map { enumCase in + let properties = enumCase.payload.map { value in + "GenerationSchema.Property(name: \"\(value.property)\", type: \(value.type).self)" + }.joined(separator: ", ") + return "(name: \"\(enumCase.name)\", properties: [\(properties)])" + }.joined(separator: ",\n ") return DeclSyntax( stringLiteral: """ @@ -1431,9 +1215,8 @@ public struct GenerableMacro: MemberMacro, ExtensionMacro { return GenerationSchema( type: Self.self, description: \(description.map { "\"\($0)\"" } ?? "\"Generated \(enumName)\""), - properties: [ - \(caseProperty), - \(valueProperty) + cases: [ + \(caseSchemas) ] ) } @@ -1490,13 +1273,14 @@ private struct EnumCaseInfo { !associatedValues.isEmpty } - var isSingleUnlabeledValue: Bool { - associatedValues.count == 1 && associatedValues[0].label == nil - } - - var isMultipleLabeledValues: Bool { - associatedValues.count > 1 - || (associatedValues.count == 1 && associatedValues[0].label != nil) + /// Each value with the property it's generated as: its label, or `value`, `value1` and so on + /// for an unlabeled one, as Foundation Models names them. + var payload: [(label: String?, property: String, type: String, isOptional: Bool)] { + associatedValues.enumerated().map { index, value in + let type = value.type + let isOptional = type.hasSuffix("?") || type.hasPrefix("Optional<") + return (value.label, value.label ?? (index == 0 ? "value" : "value\(index)"), type, isOptional) + } } } diff --git a/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift b/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift new file mode 100644 index 00000000..8f26e4d6 --- /dev/null +++ b/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift @@ -0,0 +1,96 @@ +import Foundation +import Testing + +@testable import AnyLanguageModel + +@Generable +private enum Operand: Equatable { + case name(String) + case number(Double) +} + +@Generable(description: "One thing a clause does.") +private enum Step: Equatable { + case set(field: String, to: Operand) + case add(Operand, to: String) + case pair(Int, Int) + case note(String?) + case reset +} + +/// Enums with associated values are generated the way Foundation Models generates them: each +/// case an object whose `type` names it, with its values as properties. The expected values are +/// what Foundation Models' own `@Generable` produces for the same enums. +@Suite("Generable enums with associated values") +struct GenerableEnumPayloadTests { + private func json(_ text: String) throws -> NSDictionary { + try #require(JSONSerialization.jsonObject(with: Data(text.utf8)) as? NSDictionary) + } + + @Test func casesAreObjectsNamedByType() throws { + let expected: [(Step, String)] = [ + ( + .set(field: "count", to: .number(1)), + #"{"type": "set", "field": "count", "to": {"type": "number", "value": 1}}"# + ), + ( + .add(.name("amount"), to: "count"), + #"{"type": "add", "value": {"type": "name", "value": "amount"}, "to": "count"}"# + ), + (.pair(1, 2), #"{"type": "pair", "value": 1, "value1": 2}"#), + (.reset, #"{"type": "reset"}"#), + ] + for (step, content) in expected { + #expect(try json(step.generatedContent.jsonString) == json(content)) + #expect(try Step(GeneratedContent(json: content)) == step) + } + } + + @Test func typeComesFirst() throws { + guard case .structure(_, let keys) = Step.set(field: "count", to: .number(1)).generatedContent.kind else { + Issue.record("Expected a structure") + return + } + #expect(keys == ["type", "field", "to"]) + } + + @Test func aMissingOptionalValueIsLeftOut() throws { + #expect(try json(Step.note(nil).generatedContent.jsonString) == json(#"{"type": "note"}"#)) + #expect(try Step(GeneratedContent(json: #"{"type": "note"}"#)) == .note(nil)) + #expect(try Step(GeneratedContent(json: #"{"type": "note", "value": "hi"}"#)) == .note("hi")) + } + + @Test func anUnknownTypeIsAnError() { + #expect(throws: DecodingError.self) { + try Step(GeneratedContent(json: #"{"type": "jump"}"#)) + } + } + + @Test func schemaHasAnObjectForEachCase() throws { + let encoded = try JSONEncoder().encode(Step.generationSchema) + let schema = try #require(JSONSerialization.jsonObject(with: encoded) as? [String: Any]) + let defs = try #require(schema["$defs"] as? [String: [String: Any]]) + func caseSchema(_ name: String) throws -> [String: Any] { + try #require(defs.first { $0.key.hasSuffix(".Discriminated\(name)") }?.value) + } + + let root = try #require(defs.first { $0.key.hasSuffix(".Step") }?.value) + #expect((root["anyOf"] as? [Any])?.count == 5) + + let set = try caseSchema("Set") + let properties = try #require(set["properties"] as? [String: [String: Any]]) + #expect(Set(properties.keys) == ["type", "field", "to"]) + #expect(properties["type"]?["enum"] as? [String] == ["set"]) + #expect(Set(set["required"] as? [String] ?? []) == ["type", "field", "to"]) + #expect(set["description"] as? String == "One thing a clause does.") + + let pair = try caseSchema("Pair") + #expect(Set((pair["properties"] as? [String: Any] ?? [:]).keys) == ["type", "value", "value1"]) + + let note = try caseSchema("Note") + #expect(note["required"] as? [String] == ["type"]) + + let reset = try caseSchema("Reset") + #expect(Set((reset["properties"] as? [String: Any] ?? [:]).keys) == ["type"]) + } +} diff --git a/Tests/AnyLanguageModelTests/SystemLanguageModelTests.swift b/Tests/AnyLanguageModelTests/SystemLanguageModelTests.swift index 44690e50..b9135948 100644 --- a/Tests/AnyLanguageModelTests/SystemLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/SystemLanguageModelTests.swift @@ -73,6 +73,19 @@ import Testing var titles: [String] } + @Generable + private enum CounterStep { + case add(amount: Int) + case reset + case report(String) + } + + @Generable + private struct CounterPlan { + @Guide(description: "The steps, in order.") + var steps: [CounterStep] + } + @Generable private struct SentimentAnalysis { @Guide(description: "The sentiment classification", .anyOf(["positive", "negative", "neutral"])) @@ -447,6 +460,19 @@ import Testing #expect(response.content.confidence >= 0.0 && response.content.confidence <= 1.0) } + /// An enum with associated values reaches Foundation Models as one object per case. + @available(macOS 26.0, iOS 26.0, tvOS 26.0, visionOS 26.0, *) + @Test func guidedGenerationWithEnumPayloads() async throws { + let session = LanguageModelSession(model: SystemLanguageModel.default) + + let response = try await session.respond( + to: "Plan a counter: add 2, then add 3, then report the total.", + generating: CounterPlan.self + ) + + #expect(!response.content.steps.isEmpty) + } + @available(macOS 26.0, iOS 26.0, tvOS 26.0, visionOS 26.0, *) @Test func guidedGenerationWithInstructions() async throws { let session = LanguageModelSession( From c95395bf047873784c23d7cde7d0245051e8c933 Mon Sep 17 00:00:00 2001 From: mattt Date: Tue, 13 Oct 2026 04:50:44 -0700 Subject: [PATCH 2/9] Give each enum case's schema a unique name Cases whose names differ only in their first letter, like foo and Foo, shared one Discriminated definition, so one case lost its schema. A case could also take the name of a value's nested type. A name that's taken now gets the first free numeric suffix. --- .../AnyLanguageModel/GenerationSchema.swift | 31 ++++++-- .../GenerableEnumPayloadTests.swift | 73 +++++++++++++++++++ 2 files changed, 97 insertions(+), 7 deletions(-) diff --git a/Sources/AnyLanguageModel/GenerationSchema.swift b/Sources/AnyLanguageModel/GenerationSchema.swift index 08011100..e7c29a41 100644 --- a/Sources/AnyLanguageModel/GenerationSchema.swift +++ b/Sources/AnyLanguageModel/GenerationSchema.swift @@ -448,6 +448,19 @@ public struct GenerationSchema: Equatable, Codable, CustomDebugStringConvertible var members: [Node] = [] var allDefs: [String: Node] = [:] + // Collect the values' types first, + // so that no case's definition takes the name of one of them. + for (_, properties) in cases { + for property in properties { + for (defName, defNode) in property.deps { + if let existing = allDefs[defName], existing != defNode { + fatalError("Duplicate type '\(defName)' with different structure") + } + allDefs[defName] = defNode + } + } + } + for (name, properties) in cases { var props: [String: Node] = ["type": .string(StringNode(enumChoices: [name]))] var required: Set = ["type"] @@ -456,14 +469,18 @@ public struct GenerationSchema: Equatable, Codable, CustomDebugStringConvertible if !property.isOptional { required.insert(property.name) } - for (defName, defNode) in property.deps { - if let existing = allDefs[defName], existing != defNode { - fatalError("Duplicate type '\(defName)' with different structure") - } - allDefs[defName] = defNode - } } - let caseName = "\(typeName).Discriminated\(name.prefix(1).uppercased() + name.dropFirst())" + // Name the case as Foundation Models does, after the enumeration. + // Cases whose names differ only in their first letter, like `foo` and `Foo`, + // or a value's type with the same name, + // get the first free numeric suffix, starting at 2. + let baseName = "\(typeName).Discriminated\(name.prefix(1).uppercased() + name.dropFirst())" + var caseName = baseName + var suffix = 2 + while caseName == typeName || allDefs[caseName] != nil { + caseName = "\(baseName)\(suffix)" + suffix += 1 + } allDefs[caseName] = .object( ObjectNode( description: description, diff --git a/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift b/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift index 8f26e4d6..1305ec1b 100644 --- a/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift +++ b/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift @@ -18,6 +18,22 @@ private enum Step: Equatable { case reset } +@Generable +private enum Casing: Equatable { + case foo(Int) + case Foo(String) +} + +@Generable +private enum Shelf: Equatable { + @Generable + struct DiscriminatedItem: Equatable { + var title: String + } + + case item(DiscriminatedItem) +} + /// Enums with associated values are generated the way Foundation Models generates them: each /// case an object whose `type` names it, with its values as properties. The expected values are /// what Foundation Models' own `@Generable` produces for the same enums. @@ -93,4 +109,61 @@ struct GenerableEnumPayloadTests { let reset = try caseSchema("Reset") #expect(Set((reset["properties"] as? [String: Any] ?? [:]).keys) == ["type"]) } + + /// The encoded definitions of a schema, and the definitions its root's choices refer to. + private func caseDefinitions( + _ schema: GenerationSchema + ) throws -> (defs: [String: [String: Any]], choices: [String]) { + let encoded = try JSONEncoder().encode(schema) + let object = try #require(JSONSerialization.jsonObject(with: encoded) as? [String: Any]) + let defs = try #require(object["$defs"] as? [String: [String: Any]]) + let rootName = try #require((object["$ref"] as? String)?.replacingOccurrences(of: "#/$defs/", with: "")) + let anyOf = try #require(defs[rootName]?["anyOf"] as? [[String: String]]) + let choices = anyOf.compactMap { $0["$ref"]?.replacingOccurrences(of: "#/$defs/", with: "") } + return (defs, choices) + } + + @Test func casesDifferingInCaseHaveTheirOwnSchemas() throws { + let (defs, choices) = try caseDefinitions(Casing.generationSchema) + + #expect(choices.count == 2) + #expect(Set(choices).count == 2) + #expect(choices.first?.hasSuffix(".DiscriminatedFoo") == true) + #expect(choices.last?.hasSuffix(".DiscriminatedFoo2") == true) + let types = choices.map { name in + (defs[name]?["properties"] as? [String: [String: Any]])?["type"]?["enum"] as? [String] + } + #expect(types == [["foo"], ["Foo"]]) + let values = choices.map { name in + (defs[name]?["properties"] as? [String: [String: Any]])?["value"]?["type"] as? String + } + #expect(values == ["integer", "string"]) + + // The names depend only on the declaration. + #expect(Casing.generationSchema == Casing.generationSchema) + + let expected: [(Casing, String)] = [ + (.foo(1), #"{"type":"foo","value":1}"#), + (.Foo("x"), #"{"type":"Foo","value":"x"}"#), + ] + for (casing, content) in expected { + #expect(casing.generatedContent.jsonString == content) + #expect(try Casing(GeneratedContent(json: content)) == casing) + } + } + + @Test func aCaseDoesNotTakeTheNameOfAValueType() throws { + let (defs, choices) = try caseDefinitions(Shelf.generationSchema) + let itemName = String(reflecting: Shelf.DiscriminatedItem.self) + + let caseName = try #require(choices.first) + #expect(caseName == "\(itemName)2") + let caseProperties = try #require(defs[caseName]?["properties"] as? [String: [String: Any]]) + #expect(caseProperties["value"]?["$ref"] as? String == "#/$defs/\(itemName)") + let itemProperties = try #require(defs[itemName]?["properties"] as? [String: Any]) + #expect(Set(itemProperties.keys) == ["title"]) + + let shelf = Shelf.item(.init(title: "Dune")) + #expect(try Shelf(GeneratedContent(json: shelf.generatedContent.jsonString)) == shelf) + } } From a71bfafcab35988ff0c42390e2e74ce1dd500f35 Mon Sep 17 00:00:00 2001 From: mattt Date: Tue, 13 Oct 2026 04:51:50 -0700 Subject: [PATCH 3/9] Read a null optional property as nil Foundation Models reads a null property as nil when the property is optional, and its @Generable writes a nil enum value as null. AnyLanguageModel threw a type mismatch, so such content didn't decode. --- Sources/AnyLanguageModel/GeneratedContent.swift | 6 ++++++ .../GenerableEnumPayloadTests.swift | 5 +++++ .../GeneratedContentJSONTests.swift | 10 ++++++++++ 3 files changed, 21 insertions(+) diff --git a/Sources/AnyLanguageModel/GeneratedContent.swift b/Sources/AnyLanguageModel/GeneratedContent.swift index 6c4c2fd4..c96bb0b7 100644 --- a/Sources/AnyLanguageModel/GeneratedContent.swift +++ b/Sources/AnyLanguageModel/GeneratedContent.swift @@ -349,6 +349,9 @@ public struct GeneratedContent: Sendable, Equatable, Generable, CustomDebugStrin } /// Reads an optional, concrete generable type from named property. + /// + /// Like Foundation Models, this returns `nil` + /// for a property that's missing or `null`. public func value( _ type: Value?.Type = Value?.self, forProperty property: String @@ -359,6 +362,9 @@ public struct GeneratedContent: Sendable, Equatable, Generable, CustomDebugStrin guard let value = properties[property] else { return nil } + if case .null = value.kind { + return nil + } return try Value(value) } diff --git a/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift b/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift index 1305ec1b..93b0f024 100644 --- a/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift +++ b/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift @@ -76,6 +76,11 @@ struct GenerableEnumPayloadTests { #expect(try Step(GeneratedContent(json: #"{"type": "note", "value": "hi"}"#)) == .note("hi")) } + /// Foundation Models writes a `nil` value as `null`, and reads `null` as `nil`. + @Test func aNullOptionalValueIsNil() throws { + #expect(try Step(GeneratedContent(json: #"{"type": "note", "value": null}"#)) == .note(nil)) + } + @Test func anUnknownTypeIsAnError() { #expect(throws: DecodingError.self) { try Step(GeneratedContent(json: #"{"type": "jump"}"#)) diff --git a/Tests/AnyLanguageModelTests/GeneratedContentJSONTests.swift b/Tests/AnyLanguageModelTests/GeneratedContentJSONTests.swift index 7649077d..54a511cf 100644 --- a/Tests/AnyLanguageModelTests/GeneratedContentJSONTests.swift +++ b/Tests/AnyLanguageModelTests/GeneratedContentJSONTests.swift @@ -219,6 +219,16 @@ struct GeneratedContentJSONTests { #expect(GeneratedContent(value, id: id).id == id) } + /// Like Foundation Models, reading an optional property returns `nil` for `null`. + @Test func optionalPropertyReadsNullAsNil() throws { + let content = try GeneratedContent(json: #"{"none": null, "some": "x"}"#) + #expect(try content.value(String?.self, forProperty: "none") == nil) + #expect(try content.value(GeneratedContent?.self, forProperty: "none") == nil) + #expect(try content.value(String?.self, forProperty: "missing") == nil) + #expect(try content.value(String?.self, forProperty: "some") == "x") + #expect(try content.value(GeneratedContent.self, forProperty: "none").kind == .null) + } + @Test func jsonValueRoundTrips() throws { let value: JSONValue = [ "title": "Dune", From c4a275d0e9f233cfa78e654fe2b5da51c88a7303 Mon Sep 17 00:00:00 2001 From: mattt Date: Tue, 13 Oct 2026 04:54:53 -0700 Subject: [PATCH 4/9] Write nil enum values as null, as Foundation Models does The macro left out a nil value only when its type was spelled T? or Optional, so Swift.Optional and a type alias like String? wrote null instead. Foundation Models' @Generable writes null for every nil value, whatever its spelling, so the macro now does the same and no longer checks the spelling. The schema and decoding already use overload resolution. --- .../GenerableMacro.swift | 27 ++++--- ...ableEnumPayloadFoundationModelsTests.swift | 74 +++++++++++++++++++ .../GenerableEnumPayloadTests.swift | 53 ++++++++++++- .../Shared/FoundationModelsEnumFixtures.swift | 36 +++++++++ 4 files changed, 172 insertions(+), 18 deletions(-) create mode 100644 Tests/AnyLanguageModelTests/GenerableEnumPayloadFoundationModelsTests.swift create mode 100644 Tests/AnyLanguageModelTests/Shared/FoundationModelsEnumFixtures.swift diff --git a/Sources/AnyLanguageModelMacros/GenerableMacro.swift b/Sources/AnyLanguageModelMacros/GenerableMacro.swift index 0b97b62d..8d219b24 100644 --- a/Sources/AnyLanguageModelMacros/GenerableMacro.swift +++ b/Sources/AnyLanguageModelMacros/GenerableMacro.swift @@ -1159,21 +1159,17 @@ public struct GenerableMacro: MemberMacro, ExtensionMacro { let payload = enumCase.payload let bindings = payload.indices.map { "let value\($0)" }.joined(separator: ", ") let pattern = payload.isEmpty ? ".\(enumCase.name)" : ".\(enumCase.name)(\(bindings))" - // An optional value that's nil is left out, as for a structure's properties. + // Like Foundation Models, write an optional value that's nil as null, + // whatever the spelling of its type. let assignments = payload.enumerated().map { index, value in - let assignment = """ - properties["\(value.property)"] = value\(index).generatedContent - keys.append("\(value.property)") - """ - return value.isOptional - ? "if let value\(index) {\n \(assignment)\n }" : assignment + "properties[\"\(value.property)\"] = value\(index).generatedContent" }.joined(separator: "\n ") + let keys = (["type"] + payload.map(\.property)).map { "\"\($0)\"" }.joined(separator: ", ") return """ case \(pattern): - var properties: [String: GeneratedContent] = ["type": GeneratedContent("\(enumCase.name)")] - var keys = ["type"] + \(payload.isEmpty ? "let" : "var") properties: [String: GeneratedContent] = ["type": GeneratedContent("\(enumCase.name)")] \(assignments) - return GeneratedContent(kind: .structure(properties: properties, orderedKeys: keys)) + return GeneratedContent(kind: .structure(properties: properties, orderedKeys: [\(keys)])) """ }.joined(separator: "\n ") @@ -1294,11 +1290,14 @@ private struct EnumCaseInfo { /// Each value with the property it's generated as: its label, or `value`, `value1` and so on /// for an unlabeled one, as Foundation Models names them. - var payload: [(label: String?, property: String, type: String, isOptional: Bool)] { + /// + /// The generated code doesn't depend on whether a value's type is optional: + /// overload resolution makes an optional value's property optional in the schema, + /// and reads a missing or `null` property as `nil`. + /// That covers a spelling the macro can't resolve, like a type alias. + var payload: [(label: String?, property: String, type: String)] { associatedValues.enumerated().map { index, value in - let type = value.type - let isOptional = type.hasSuffix("?") || type.hasPrefix("Optional<") - return (value.label, value.label ?? (index == 0 ? "value" : "value\(index)"), type, isOptional) + (value.label, value.label ?? (index == 0 ? "value" : "value\(index)"), value.type) } } } diff --git a/Tests/AnyLanguageModelTests/GenerableEnumPayloadFoundationModelsTests.swift b/Tests/AnyLanguageModelTests/GenerableEnumPayloadFoundationModelsTests.swift new file mode 100644 index 00000000..8b50e76f --- /dev/null +++ b/Tests/AnyLanguageModelTests/GenerableEnumPayloadFoundationModelsTests.swift @@ -0,0 +1,74 @@ +import Foundation +import Testing + +@testable import AnyLanguageModel + +#if canImport(FoundationModels) && compiler(>=6.4) && !os(tvOS) + import FoundationModels + + private let isFoundationModelsAvailable: Bool = { + if #available(macOS 26.0, iOS 26.0, watchOS 27.0, visionOS 26.0, *) { + return true + } + return false + }() + + /// Compares enums with associated values in AnyLanguageModel + /// with the same enums in Foundation Models, + /// whose `@Generable` macro comes from the SDK that builds the tests. + @Suite("Generable enums with associated values in Foundation Models", .enabled(if: isFoundationModelsAvailable)) + struct GenerableEnumPayloadFoundationModelsTests { + @available(macOS 26.0, iOS 26.0, watchOS 27.0, visionOS 26.0, *) + private typealias Native = FoundationModelsEnums + + @available(macOS 26.0, iOS 26.0, watchOS 27.0, visionOS 26.0, *) + private var steps: [(Step, Native.Step)] { + [ + (.set(field: "count", to: .number(1)), .set(field: "count", to: .number(1))), + (.add(.name("amount"), to: "count"), .add(.name("amount"), to: "count")), + (.pair(1, 2), .pair(1, 2)), + (.note("hi"), .note("hi")), + (.note(nil), .note(nil)), + (.reset, .reset), + ] + } + + @available(macOS 26.0, iOS 26.0, watchOS 27.0, visionOS 26.0, *) + private var annotations: [(Annotation, Native.Annotation)] { + [ + (.alias(nil), .alias(nil)), + (.generic(nil), .generic(nil)), + (.qualified(nil), .qualified(nil)), + (.alias("hi"), .alias("hi")), + (.generic(3), .generic(3)), + ] + } + + /// Foundation Models writes a space after each colon and comma. + private func compacted(_ json: String) -> String { + json.replacingOccurrences(of: ": ", with: ":").replacingOccurrences(of: ", ", with: ",") + } + + @available(macOS 26.0, iOS 26.0, watchOS 27.0, visionOS 26.0, *) + @Test func contentMatchesFoundationModels() { + for (step, native) in steps { + #expect(step.generatedContent.jsonString == compacted(native.generatedContent.jsonString)) + } + for (annotation, native) in annotations { + #expect(annotation.generatedContent.jsonString == compacted(native.generatedContent.jsonString)) + } + } + + @available(macOS 26.0, iOS 26.0, watchOS 27.0, visionOS 26.0, *) + @Test func foundationModelsContentDecodes() throws { + for (step, native) in steps { + let content = try AnyLanguageModel.GeneratedContent(json: native.generatedContent.jsonString) + #expect(try Step(content) == step) + } + for (annotation, native) in annotations { + let content = try AnyLanguageModel.GeneratedContent(json: native.generatedContent.jsonString) + #expect(try Annotation(content) == annotation) + } + } + } +#endif diff --git a/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift b/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift index 93b0f024..8fa2d2a3 100644 --- a/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift +++ b/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift @@ -3,14 +3,17 @@ import Testing @testable import AnyLanguageModel +// `GenerableEnumPayloadFoundationModelsTests` compares these enums +// with the same enums in Foundation Models. + @Generable -private enum Operand: Equatable { +enum Operand: Equatable { case name(String) case number(Double) } @Generable(description: "One thing a clause does.") -private enum Step: Equatable { +enum Step: Equatable { case set(field: String, to: Operand) case add(Operand, to: String) case pair(Int, Int) @@ -18,6 +21,17 @@ private enum Step: Equatable { case reset } +typealias MaybeText = String? + +/// Optional values, with each spelling of an optional type. +@Generable +enum Annotation: Equatable { + case note(String?) + case alias(MaybeText) + case generic(Optional) + case qualified(Swift.Optional) +} + @Generable private enum Casing: Equatable { case foo(Int) @@ -70,8 +84,7 @@ struct GenerableEnumPayloadTests { #expect(keys == ["type", "field", "to"]) } - @Test func aMissingOptionalValueIsLeftOut() throws { - #expect(try json(Step.note(nil).generatedContent.jsonString) == json(#"{"type": "note"}"#)) + @Test func aMissingOptionalValueIsNil() throws { #expect(try Step(GeneratedContent(json: #"{"type": "note"}"#)) == .note(nil)) #expect(try Step(GeneratedContent(json: #"{"type": "note", "value": "hi"}"#)) == .note("hi")) } @@ -81,6 +94,38 @@ struct GenerableEnumPayloadTests { #expect(try Step(GeneratedContent(json: #"{"type": "note", "value": null}"#)) == .note(nil)) } + /// The macro can't see through a type alias, + /// so nothing it generates depends on how an optional type is spelled. + @Test func optionalValuesDoNotDependOnSpelling() throws { + let expected: [(Annotation, String)] = [ + (.note(nil), #"{"type":"note","value":null}"#), + (.alias(nil), #"{"type":"alias","value":null}"#), + (.generic(nil), #"{"type":"generic","value":null}"#), + (.qualified(nil), #"{"type":"qualified","value":null}"#), + (.note("a"), #"{"type":"note","value":"a"}"#), + (.alias("b"), #"{"type":"alias","value":"b"}"#), + (.generic(2), #"{"type":"generic","value":2}"#), + (.qualified("c"), #"{"type":"qualified","value":"c"}"#), + ] + for (annotation, content) in expected { + #expect(annotation.generatedContent.jsonString == content) + #expect(try Annotation(GeneratedContent(json: content)) == annotation) + } + + let missing: [(String, Annotation)] = [ + ("note", .note(nil)), ("alias", .alias(nil)), ("generic", .generic(nil)), ("qualified", .qualified(nil)), + ] + for (type, annotation) in missing { + #expect(try Annotation(GeneratedContent(json: #"{"type":"\#(type)"}"#)) == annotation) + } + + let (defs, choices) = try caseDefinitions(Annotation.generationSchema) + #expect(choices.count == 4) + for name in choices { + #expect(defs[name]?["required"] as? [String] == ["type"]) + } + } + @Test func anUnknownTypeIsAnError() { #expect(throws: DecodingError.self) { try Step(GeneratedContent(json: #"{"type": "jump"}"#)) diff --git a/Tests/AnyLanguageModelTests/Shared/FoundationModelsEnumFixtures.swift b/Tests/AnyLanguageModelTests/Shared/FoundationModelsEnumFixtures.swift new file mode 100644 index 00000000..0f8cb919 --- /dev/null +++ b/Tests/AnyLanguageModelTests/Shared/FoundationModelsEnumFixtures.swift @@ -0,0 +1,36 @@ +// Foundation Models' own `@Generable` enums, +// declared like the enums in `GenerableEnumPayloadTests`. +// This file doesn't import AnyLanguageModel, +// whose macro expansions would make names like `GeneratedContent` ambiguous. + +#if canImport(FoundationModels) && compiler(>=6.4) && !os(tvOS) + import FoundationModels + + @available(macOS 26.0, iOS 26.0, watchOS 27.0, visionOS 26.0, *) + enum FoundationModelsEnums { + @Generable + enum Operand: Equatable { + case name(String) + case number(Double) + } + + @Generable(description: "One thing a clause does.") + enum Step: Equatable { + case set(field: String, to: Operand) + case add(Operand, to: String) + case pair(Int, Int) + case note(String?) + case reset + } + + typealias MaybeText = String? + + @Generable + enum Annotation: Equatable { + case note(String?) + case alias(MaybeText) + case generic(Optional) + case qualified(Swift.Optional) + } + } +#endif From 622e3a3c07365971f797ba8312917d8982a85e26 Mon Sep 17 00:00:00 2001 From: mattt Date: Tue, 13 Oct 2026 04:57:04 -0700 Subject: [PATCH 5/9] Describe each enum case's type as a constant, as Foundation Models does Foundation Models' @Generable guides an enum case's type with .constant, which encodes as const. Since #304, providers whose schemas don't support const already get a one-choice enum, so the schema can use a constant everywhere and still work with OpenAI, Open Responses, Gemini, and Ollama. --- .../AnyLanguageModel/GenerationSchema.swift | 4 +- ...ableEnumPayloadFoundationModelsTests.swift | 41 +++++++++++++++++++ .../GenerableEnumPayloadTests.swift | 23 +++++++++-- 3 files changed, 64 insertions(+), 4 deletions(-) diff --git a/Sources/AnyLanguageModel/GenerationSchema.swift b/Sources/AnyLanguageModel/GenerationSchema.swift index e7c29a41..eace2cf6 100644 --- a/Sources/AnyLanguageModel/GenerationSchema.swift +++ b/Sources/AnyLanguageModel/GenerationSchema.swift @@ -462,7 +462,9 @@ public struct GenerationSchema: Equatable, Codable, CustomDebugStringConvertible } for (name, properties) in cases { - var props: [String: Node] = ["type": .string(StringNode(enumChoices: [name]))] + // Like Foundation Models, `type` is a constant. + // Providers whose schemas don't support `const` get a one-choice `enum`. + var props: [String: Node] = ["type": .string(StringNode(enumChoices: [name], isConstant: true))] var required: Set = ["type"] for property in properties { props[property.name] = property.node diff --git a/Tests/AnyLanguageModelTests/GenerableEnumPayloadFoundationModelsTests.swift b/Tests/AnyLanguageModelTests/GenerableEnumPayloadFoundationModelsTests.swift index 8b50e76f..5d4813d2 100644 --- a/Tests/AnyLanguageModelTests/GenerableEnumPayloadFoundationModelsTests.swift +++ b/Tests/AnyLanguageModelTests/GenerableEnumPayloadFoundationModelsTests.swift @@ -59,6 +59,47 @@ import Testing } } + /// The encoded definitions of each case in a schema, by the name in its `type`. + @available(macOS 26.0, iOS 26.0, watchOS 27.0, visionOS 26.0, *) + private func caseDefinitions(_ schema: FoundationModels.GenerationSchema) throws -> [String: NSDictionary] { + let encoder = JSONEncoder() + encoder.outputFormatting = .sortedKeys + let object = try JSONSerialization.jsonObject(with: encoder.encode(schema)) as? [String: Any] + let defs = try #require(object?["$defs"] as? [String: NSDictionary]) + var cases: [String: NSDictionary] = [:] + for definition in defs.values { + guard let type = (definition["properties"] as? [String: NSDictionary])?["type"], + let name = type["const"] as? String ?? (type["enum"] as? [String])?.first + else { continue } + cases[name] = definition + } + return cases + } + + /// Converted for the system model, each case's schema encodes as Foundation Models' own does, + /// apart from the names of definitions in `$ref`s, the default descriptions, + /// and the order of properties, which the conversion doesn't keep. + @available(macOS 26.0, iOS 26.0, watchOS 27.0, visionOS 26.0, *) + @Test func caseSchemasMatchFoundationModels() throws { + let converted = try caseDefinitions(FoundationModels.GenerationSchema(Step.generationSchema)) + let native = try caseDefinitions(Native.Step.generationSchema) + + #expect(Set(converted.keys) == Set(native.keys)) + #expect(Set(native.keys).isSuperset(of: ["set", "add", "pair", "note", "reset", "name", "number"])) + for (name, nativeCase) in native { + let convertedCase = try #require(converted[name]) + #expect( + Set(convertedCase["required"] as? [String] ?? []) == Set(nativeCase["required"] as? [String] ?? []) + ) + let properties = try #require(convertedCase["properties"] as? [String: NSDictionary]) + let nativeProperties = try #require(nativeCase["properties"] as? [String: NSDictionary]) + #expect(Set(properties.keys) == Set(nativeProperties.keys)) + for (key, nativeProperty) in nativeProperties where nativeProperty["$ref"] == nil { + #expect(properties[key] == nativeProperty, "\(name).\(key)") + } + } + } + @available(macOS 26.0, iOS 26.0, watchOS 27.0, visionOS 26.0, *) @Test func foundationModelsContentDecodes() throws { for (step, native) in steps { diff --git a/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift b/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift index 8fa2d2a3..6c453068 100644 --- a/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift +++ b/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift @@ -146,7 +146,7 @@ struct GenerableEnumPayloadTests { let set = try caseSchema("Set") let properties = try #require(set["properties"] as? [String: [String: Any]]) #expect(Set(properties.keys) == ["type", "field", "to"]) - #expect(properties["type"]?["enum"] as? [String] == ["set"]) + #expect(properties["type"]?["const"] as? String == "set") #expect(Set(set["required"] as? [String] ?? []) == ["type", "field", "to"]) #expect(set["description"] as? String == "One thing a clause does.") @@ -160,6 +160,23 @@ struct GenerableEnumPayloadTests { #expect(Set((reset["properties"] as? [String: Any] ?? [:]).keys) == ["type"]) } + /// Like Foundation Models, a case's `type` is a constant, + /// which providers without `const` get as a one-choice `enum`. + @Test func typeIsAConstant() throws { + let inlined = try JSONValue(Step.generationSchema.inlinedJSONSchema()) + let asEnum = try JSONValue(Step.generationSchema.inlinedJSONSchema(constantsAsEnums: true)) + guard case .array(let cases)? = inlined.objectValue?["anyOf"], + case .array(let enumCases)? = asEnum.objectValue?["anyOf"] + else { + Issue.record("Expected a union of cases") + return + } + let type = cases.first?.objectValue?["properties"]?.objectValue?["type"] + #expect(type == ["type": "string", "const": "set"]) + let enumType = enumCases.first?.objectValue?["properties"]?.objectValue?["type"] + #expect(enumType == ["type": "string", "enum": ["set"]]) + } + /// The encoded definitions of a schema, and the definitions its root's choices refer to. private func caseDefinitions( _ schema: GenerationSchema @@ -181,9 +198,9 @@ struct GenerableEnumPayloadTests { #expect(choices.first?.hasSuffix(".DiscriminatedFoo") == true) #expect(choices.last?.hasSuffix(".DiscriminatedFoo2") == true) let types = choices.map { name in - (defs[name]?["properties"] as? [String: [String: Any]])?["type"]?["enum"] as? [String] + (defs[name]?["properties"] as? [String: [String: Any]])?["type"]?["const"] as? String } - #expect(types == [["foo"], ["Foo"]]) + #expect(types == ["foo", "Foo"]) let values = choices.map { name in (defs[name]?["properties"] as? [String: [String: Any]])?["value"]?["type"] as? String } From fcd497787d5875a37ae6a0bc6e183cf480752586 Mon Sep 17 00:00:00 2001 From: mattt Date: Tue, 13 Oct 2026 05:01:36 -0700 Subject: [PATCH 6/9] Wrap response schemas without an object root for providers that need one An enum whose cases have associated values now has an anyOf at the root of its schema, as in Foundation Models. OpenAI's structured outputs reject a root that isn't an object or that uses anyOf, and Open Responses follows them, so OpenAI and Open Responses get any other root as the response property of an object. Gemini's response schemas can be objects, primitives, or arrays, so Gemini gets a union root wrapped the same way. Each provider unwraps the content it returns, streamed or not. Content without the wrapper, from a provider that doesn't apply the schema, is used as sent. GenerationSchema itself keeps the Foundation Models form. --- .../AnyLanguageModel/GenerationSchema.swift | 2 +- .../LanguageModelSession.swift | 12 +- .../Models/GeminiLanguageModel.swift | 12 +- .../Models/OpenAILanguageModel.swift | 21 ++-- .../Models/OpenResponsesLanguageModel.swift | 12 +- .../Shared/ProviderResponseSchema.swift | 95 ++++++++++++++ .../Shared/ResponseUsage.swift | 10 +- .../Shared/StreamingResponseState.swift | 13 +- .../ProviderSchemaResponseTests.swift | 116 ++++++++++++++++++ 9 files changed, 267 insertions(+), 26 deletions(-) create mode 100644 Sources/AnyLanguageModel/Shared/ProviderResponseSchema.swift diff --git a/Sources/AnyLanguageModel/GenerationSchema.swift b/Sources/AnyLanguageModel/GenerationSchema.swift index eace2cf6..41a13d13 100644 --- a/Sources/AnyLanguageModel/GenerationSchema.swift +++ b/Sources/AnyLanguageModel/GenerationSchema.swift @@ -603,7 +603,7 @@ public struct GenerationSchema: Equatable, Codable, CustomDebugStringConvertible self.defs = allDefs } - private init(root: Node, defs: [String: Node]) { + init(root: Node, defs: [String: Node]) { self.root = root self.defs = defs } diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index 4f08efc6..38a7dd1a 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/Sources/AnyLanguageModel/LanguageModelSession.swift @@ -863,7 +863,7 @@ extension LanguageModelSession { extension LanguageModelSession { /// Generates content that conforms to the supplied schema. - /// Provider restrictions still apply; OpenAI strict mode requires an object root. + /// Provider restrictions on schemas still apply. @discardableResult nonisolated public func respond( to prompt: Prompt, @@ -883,7 +883,7 @@ extension LanguageModelSession { } /// Generates content that conforms to the supplied schema. - /// Provider restrictions still apply; OpenAI strict mode requires an object root. + /// Provider restrictions on schemas still apply. @discardableResult nonisolated public func respond( to prompt: String, @@ -900,7 +900,7 @@ extension LanguageModelSession { } /// Generates content that conforms to the supplied schema. - /// Provider restrictions still apply; OpenAI strict mode requires an object root. + /// Provider restrictions on schemas still apply. @discardableResult nonisolated public func respond( schema: GenerationSchema, @@ -917,7 +917,7 @@ extension LanguageModelSession { } /// Streams content that conforms to the supplied schema. - /// Provider restrictions still apply; OpenAI strict mode requires an object root. + /// Provider restrictions on schemas still apply. nonisolated public func streamResponse( to prompt: Prompt, schema: GenerationSchema, @@ -948,7 +948,7 @@ extension LanguageModelSession { } /// Streams content that conforms to the supplied schema. - /// Provider restrictions still apply; OpenAI strict mode requires an object root. + /// Provider restrictions on schemas still apply. nonisolated public func streamResponse( to prompt: String, schema: GenerationSchema, @@ -964,7 +964,7 @@ extension LanguageModelSession { } /// Streams content that conforms to the supplied schema. - /// Provider restrictions still apply; OpenAI strict mode requires an object root. + /// Provider restrictions on schemas still apply. nonisolated public func streamResponse( schema: GenerationSchema, includeSchemaInPrompt: Bool = true, diff --git a/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift b/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift index 5fc12db9..5be74601 100644 --- a/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift @@ -223,6 +223,8 @@ public struct GeminiLanguageModel: LanguageModel { let effectiveThinking = customOptions?.thinking ?? .disabled let effectiveServerTools = customOptions?.serverTools ?? [] let effectiveJsonMode = customOptions?.jsonMode + // Gemini's response schemas can be objects, primitives, or arrays, but not unions. + let responseSchema = ProviderResponseSchema(schema, generating: type, wrapping: .unionRoots) let url = baseURL @@ -255,7 +257,7 @@ public struct GeminiLanguageModel: LanguageModel { contents: requestTranscript.toGeminiContent(), tools: geminiTools, generating: type, - schema: schema, + schema: responseSchema.schema, options: options, thinking: effectiveThinking, jsonMode: effectiveJsonMode @@ -340,7 +342,7 @@ public struct GeminiLanguageModel: LanguageModel { ) } - let generatedContent = try GeneratedContent(json: text) + let generatedContent = try responseSchema.content(json: text) let content = try type.init(generatedContent) return LanguageModelSession.Response( content: content, @@ -400,6 +402,8 @@ public struct GeminiLanguageModel: LanguageModel { let effectiveThinking = customOptions?.thinking ?? .disabled let effectiveServerTools = customOptions?.serverTools ?? [] let effectiveJsonMode = customOptions?.jsonMode + // Gemini's response schemas can be objects, primitives, or arrays, but not unions. + let responseSchema = ProviderResponseSchema(schema, generating: type, wrapping: .unionRoots) var streamURL = baseURL @@ -414,7 +418,7 @@ public struct GeminiLanguageModel: LanguageModel { do { let headers = buildHeaders() var inFlightEntries: [Transcript.Entry] = [] - var state = StreamingResponseState() + var state = StreamingResponseState(wrapperProperty: responseSchema.wrapperProperty) var toolRounds = ToolRoundLimit(provider: "Gemini") while true { try Task.checkCancellation() @@ -431,7 +435,7 @@ public struct GeminiLanguageModel: LanguageModel { contents: requestTranscript.toGeminiContent(), tools: geminiTools, generating: type, - schema: schema, + schema: responseSchema.schema, options: options, thinking: effectiveThinking, jsonMode: effectiveJsonMode diff --git a/Sources/AnyLanguageModel/Models/OpenAILanguageModel.swift b/Sources/AnyLanguageModel/Models/OpenAILanguageModel.swift index 6f17d3da..6952ab62 100644 --- a/Sources/AnyLanguageModel/Models/OpenAILanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/OpenAILanguageModel.swift @@ -492,7 +492,8 @@ public struct OpenAILanguageModel: LanguageModel { options: GenerationOptions, session: LanguageModelSession ) async throws -> LanguageModelSession.Response where Content: Generable { - + // Structured outputs need an object root. + let responseSchema = ProviderResponseSchema(schema, generating: type, wrapping: .nonObjectRoots) var entries: [Transcript.Entry] = [] var usage = ReportedUsage() var text = "" @@ -513,7 +514,7 @@ public struct OpenAILanguageModel: LanguageModel { messages: messages, tools: tools, generating: type, - schema: schema, + schema: responseSchema.schema, options: options, stream: false ) @@ -602,7 +603,7 @@ public struct OpenAILanguageModel: LanguageModel { ) } - let generatedContent = try GeneratedContent(json: text) + let generatedContent = try responseSchema.content(json: text) let content = try type.init(generatedContent) return LanguageModelSession.Response( content: content, @@ -618,6 +619,8 @@ public struct OpenAILanguageModel: LanguageModel { options: GenerationOptions, session: LanguageModelSession ) async throws -> LanguageModelSession.Response where Content: Generable { + // Structured outputs need an object root. + let responseSchema = ProviderResponseSchema(schema, generating: type, wrapping: .nonObjectRoots) var entries: [Transcript.Entry] = [] var usage = ReportedUsage() var text = "" @@ -641,7 +644,7 @@ public struct OpenAILanguageModel: LanguageModel { messages: messages, tools: tools, generating: type, - schema: schema, + schema: responseSchema.schema, options: options, stream: false ) @@ -722,7 +725,7 @@ public struct OpenAILanguageModel: LanguageModel { } if let jsonString = extractJSONFromOutput(lastOutput) { - let generatedContent = try GeneratedContent(json: jsonString) + let generatedContent = try responseSchema.content(json: jsonString) let content = try type.init(generatedContent) return LanguageModelSession.Response( content: content, @@ -776,12 +779,14 @@ public struct OpenAILanguageModel: LanguageModel { includeSchemaInPrompt: Bool, options: GenerationOptions ) -> sending LanguageModelSession.ResponseStream where Content: Generable { + // Structured outputs need an object root. + let responseSchema = ProviderResponseSchema(schema, generating: type, wrapping: .nonObjectRoots) let stream = AsyncThrowingStream.Snapshot, any Error> { continuation in let task = Task { do { var inFlightMessages: [OpenAIMessage] = [] - var state = StreamingResponseState() + var state = StreamingResponseState(wrapperProperty: responseSchema.wrapperProperty) var toolRounds = ToolRoundLimit(provider: "OpenAI") while true { try Task.checkCancellation() @@ -800,7 +805,7 @@ public struct OpenAILanguageModel: LanguageModel { messages: messages, tools: tools, generating: type, - schema: schema, + schema: responseSchema.schema, options: options, stream: true ) @@ -811,7 +816,7 @@ public struct OpenAILanguageModel: LanguageModel { messages: messages, tools: tools, generating: type, - schema: schema, + schema: responseSchema.schema, options: options, stream: true, includeUsage: baseURL.host == Self.defaultBaseURL.host diff --git a/Sources/AnyLanguageModel/Models/OpenResponsesLanguageModel.swift b/Sources/AnyLanguageModel/Models/OpenResponsesLanguageModel.swift index 7eb9992d..395b40b8 100644 --- a/Sources/AnyLanguageModel/Models/OpenResponsesLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/OpenResponsesLanguageModel.swift @@ -486,12 +486,14 @@ public struct OpenResponsesLanguageModel: LanguageModel { options: GenerationOptions ) -> sending LanguageModelSession.ResponseStream where Content: Generable { let url = baseURL.appendingPathComponent("responses") + // Open Responses follows OpenAI's structured outputs, which need an object root. + let responseSchema = ProviderResponseSchema(schema, generating: type, wrapping: .nonObjectRoots) let stream = AsyncThrowingStream.Snapshot, any Error> { continuation in let task = Task { do { var inFlightMessages: [OpenResponsesMessage] = [] - var state = StreamingResponseState() + var state = StreamingResponseState(wrapperProperty: responseSchema.wrapperProperty) var toolRounds = ToolRoundLimit(provider: "Open Responses") while true { try Task.checkCancellation() @@ -505,7 +507,7 @@ public struct OpenResponsesLanguageModel: LanguageModel { messages: messages, tools: tools, generating: type, - schema: schema, + schema: responseSchema.schema, options: options, stream: true ) @@ -593,6 +595,8 @@ public struct OpenResponsesLanguageModel: LanguageModel { options: GenerationOptions, session: LanguageModelSession ) async throws -> LanguageModelSession.Response where Content: Generable { + // Open Responses follows OpenAI's structured outputs, which need an object root. + let responseSchema = ProviderResponseSchema(schema, generating: type, wrapping: .nonObjectRoots) var entries: [Transcript.Entry] = [] var usage = ReportedUsage() var text = "" @@ -614,7 +618,7 @@ public struct OpenResponsesLanguageModel: LanguageModel { messages: messages, tools: tools, generating: type, - schema: schema, + schema: responseSchema.schema, options: options, stream: false ) @@ -689,7 +693,7 @@ public struct OpenResponsesLanguageModel: LanguageModel { ) } if let jsonString = extractJSONFromOutput(lastOutput) { - let generatedContent = try GeneratedContent(json: jsonString) + let generatedContent = try responseSchema.content(json: jsonString) let content = try type.init(generatedContent) return LanguageModelSession.Response( content: content, diff --git a/Sources/AnyLanguageModel/Shared/ProviderResponseSchema.swift b/Sources/AnyLanguageModel/Shared/ProviderResponseSchema.swift new file mode 100644 index 00000000..30f5bf13 --- /dev/null +++ b/Sources/AnyLanguageModel/Shared/ProviderResponseSchema.swift @@ -0,0 +1,95 @@ +/// A response schema in the form that a provider accepts. +/// +/// Some providers need an object at the root of a response schema. +/// OpenAI's structured outputs, for example, +/// reject a root that isn't an object or that uses `anyOf`. +/// For those providers, a schema with another root, +/// such as an enumeration whose cases have associated values or an array, +/// becomes the one required property of an object, +/// and the response is unwrapped before it's decoded. +/// +/// Wrapping happens only here, at the provider boundary, +/// so `GenerationSchema` keeps the form that Foundation Models gives it. +struct ProviderResponseSchema { + /// The roots that a provider doesn't accept. + enum Wrapping { + /// Any root that isn't an object. + case nonObjectRoots + /// A root that's a union of schemas. + case unionRoots + } + + /// The name of the property that holds a wrapped response. + static let wrapperProperty = "response" + + /// The schema to send to the provider. + let schema: GenerationSchema + + /// The property that holds the response, if the schema wraps it. + let wrapperProperty: String? + + /// Creates the response schema to send for content of the given type. + /// + /// String content sends no schema, so it's never wrapped. + init(_ schema: GenerationSchema, generating type: Content.Type, wrapping: Wrapping) { + let wraps: Bool + switch (wrapping, schema.resolvedRoot) { + case (_, .object?): + wraps = false + case (.nonObjectRoots, _), (.unionRoots, .anyOf?): + wraps = true + case (.unionRoots, _): + wraps = false + } + guard wraps, type != String.self else { + self.schema = schema + self.wrapperProperty = nil + return + } + self.schema = schema.wrappingRoot(inProperty: Self.wrapperProperty) + self.wrapperProperty = Self.wrapperProperty + } + + /// Returns the response content in JSON text, unwrapped if the schema wraps it. + func content(json text: String) throws -> GeneratedContent { + Self.unwrapping(try GeneratedContent(json: text), fromProperty: wrapperProperty) + } + + /// Returns the value of the wrapper property + /// if the content is an object with only that property. + /// + /// Other content is returned as the provider sent it, + /// so a provider that doesn't apply the response schema can return content without the wrapper. + static func unwrapping(_ content: GeneratedContent, fromProperty property: String?) -> GeneratedContent { + guard let property, + case .structure(let properties, _) = content.kind, + properties.count == 1, + let wrapped = properties[property] + else { return content } + return wrapped + } +} + +extension GenerationSchema { + /// The root, after following references. + var resolvedRoot: Node? { + var node = root + var seen: Set = [] + while case .ref(let name) = node { + guard seen.insert(name).inserted, let definition = defs[name] else { return nil } + node = definition + } + return node + } + + /// Returns the schema with its root as the one required property of an object. + func wrappingRoot(inProperty property: String) -> GenerationSchema { + let object = ObjectNode( + description: nil, + properties: [property: root], + required: [property], + propertyOrder: [property] + ) + return GenerationSchema(root: .object(object), defs: defs) + } +} diff --git a/Sources/AnyLanguageModel/Shared/ResponseUsage.swift b/Sources/AnyLanguageModel/Shared/ResponseUsage.swift index 69f192e0..60dba187 100644 --- a/Sources/AnyLanguageModel/Shared/ResponseUsage.swift +++ b/Sources/AnyLanguageModel/Shared/ResponseUsage.swift @@ -117,7 +117,10 @@ struct ResponsesUsage: Decodable, Sendable { extension LanguageModelSession.ResponseStream.Snapshot { /// Builds a snapshot from accumulated text, /// including metadata-only updates. - init?(text: String, usage: LanguageModelSession.Usage) { + /// + /// - Parameter wrapperProperty: The property that holds the response, + /// if the provider's response schema wraps it. + init?(text: String, usage: LanguageModelSession.Usage, wrapperProperty: String? = nil) { if Content.self == String.self { self.init( content: (text as! Content).asPartiallyGenerated(), @@ -125,7 +128,10 @@ extension LanguageModelSession.ResponseStream.Snapshot { usage: usage ) } else { - let raw = (try? GeneratedContent(json: text)) ?? GeneratedContent(text) + let raw = ProviderResponseSchema.unwrapping( + (try? GeneratedContent(json: text)) ?? GeneratedContent(text), + fromProperty: wrapperProperty + ) guard let content = try? Content(raw) else { return nil } self.init(content: content.asPartiallyGenerated(), rawContent: raw, usage: usage) } diff --git a/Sources/AnyLanguageModel/Shared/StreamingResponseState.swift b/Sources/AnyLanguageModel/Shared/StreamingResponseState.swift index 470e86fd..7d82d076 100644 --- a/Sources/AnyLanguageModel/Shared/StreamingResponseState.swift +++ b/Sources/AnyLanguageModel/Shared/StreamingResponseState.swift @@ -6,6 +6,13 @@ struct StreamingResponseState { var usage = ReportedUsage() private var completedUsage = LanguageModelSession.Usage.zero + /// The property that holds the response, if the provider's response schema wraps it. + let wrapperProperty: String? + + init(wrapperProperty: String? = nil) { + self.wrapperProperty = wrapperProperty + } + /// The text of completed rounds, for string content. /// /// Like the MLX and llama.cpp providers, @@ -27,7 +34,11 @@ struct StreamingResponseState { func snapshot(providerMetadata: [String: String]? = nil) -> LanguageModelSession.ResponseStream.Snapshot? { guard - var snapshot = LanguageModelSession.ResponseStream.Snapshot(text: responseText, usage: totalUsage) + var snapshot = LanguageModelSession.ResponseStream.Snapshot( + text: responseText, + usage: totalUsage, + wrapperProperty: wrapperProperty + ) else { return nil } snapshot.transcriptEntries = ArraySlice(entries) snapshot.providerMetadata = providerMetadata diff --git a/Tests/AnyLanguageModelTests/ProviderSchemaResponseTests.swift b/Tests/AnyLanguageModelTests/ProviderSchemaResponseTests.swift index c1834dcd..36633416 100644 --- a/Tests/AnyLanguageModelTests/ProviderSchemaResponseTests.swift +++ b/Tests/AnyLanguageModelTests/ProviderSchemaResponseTests.swift @@ -250,6 +250,122 @@ import Testing } } + extension ProviderSchemaResponseTests { + /// The text of the response that the session's transcript records last. + private func lastResponseText(in session: LanguageModelSession) -> String? { + guard case .response(let response)? = session.transcript.last, + case .text(let text)? = response.segments.first + else { return nil } + return text.content + } + + /// The JSON value in a string, which compares without regard to the order of keys. + private func jsonValue(_ text: String?) throws -> JSONValue? { + try text.map { try JSONDecoder().decode(JSONValue.self, from: Data($0.utf8)) } + } + + /// An enumeration whose cases have associated values has a union at the root of its schema. + /// OpenAI and Open Responses need an object root, and Gemini's response schemas have no unions, + /// so they get the union as the `response` property of an object, + /// and the content they return is unwrapped. + /// Anthropic and Ollama get the union itself. + @Test(arguments: Provider.allCases, [false, true]) + func unionRootReachesRequest(_ provider: Provider, _ streaming: Bool) async throws { + SchemaURLProtocol.reset() + defer { SchemaURLProtocol.reset() } + let wraps = provider != .anthropic && provider != .ollama + let operand = #"{"type":"number","value":1}"# + let text = wraps ? #"{"response":\#(operand)}"# : operand + SchemaURLProtocol.enqueue( + json: try streaming ? provider.stream(text: text) : json(provider.response(text: text)) + ) + let session = provider.session() + let response = + try await streaming + ? session.streamResponse(to: "One", generating: Operand.self, includeSchemaInPrompt: false).collect() + : session.respond(to: "One", generating: Operand.self, includeSchemaInPrompt: false) + #expect(response.content == .number(1)) + #expect(try jsonValue(response.rawContent.jsonString) == jsonValue(operand)) + #expect(try jsonValue(lastResponseText(in: session)) == jsonValue(operand)) + + let data = try #require(SchemaURLProtocol.recordedBodies.first) + let body = try #require(JSONSerialization.jsonObject(with: data) as? [String: Any]) + let sent = try provider.requestSchema(body) + #expect(sent["$defs"] == nil) + var union = sent + if wraps { + #expect((sent["type"] as? String)?.lowercased() == "object") + #expect(sent["required"] as? [String] == ["response"]) + let properties = try #require(sent["properties"] as? [String: [String: Any]]) + #expect(Set(properties.keys) == ["response"]) + union = try #require(properties["response"]) + if provider != .gemini { + #expect(sent["additionalProperties"] as? Bool == false) + } + } + let cases = try #require(union["anyOf"] as? [[String: Any]]) + let types = cases.map { ($0["properties"] as? [String: [String: Any]])?["type"] } + #expect(types.count == 2) + for type in types { + #expect(type?["const"] != nil || type?["enum"] != nil) + } + } + + /// OpenAI and Open Responses need an object root, so an array is wrapped too. + /// Gemini, Anthropic, and Ollama accept an array at the root. + @Test(arguments: Provider.allCases, [false, true]) + func arrayRootIsWrappedOnlyForObjectRootProviders(_ provider: Provider, _ streaming: Bool) async throws { + SchemaURLProtocol.reset() + defer { SchemaURLProtocol.reset() } + let wraps = provider == .chat || provider == .responses || provider == .openResponses + let array = #"["a","b"]"# + let text = wraps ? #"{"response":\#(array)}"# : array + SchemaURLProtocol.enqueue( + json: try streaming ? provider.stream(text: text) : json(provider.response(text: text)) + ) + let session = provider.session() + let response = + try await streaming + ? session.streamResponse(to: "Two", generating: [String].self, includeSchemaInPrompt: false).collect() + : session.respond(to: "Two", generating: [String].self, includeSchemaInPrompt: false) + #expect(response.content == ["a", "b"]) + #expect(try jsonValue(lastResponseText(in: session)) == jsonValue(array)) + + let data = try #require(SchemaURLProtocol.recordedBodies.first) + let body = try #require(JSONSerialization.jsonObject(with: data) as? [String: Any]) + let sent = try provider.requestSchema(body) + if wraps { + #expect(sent["type"] as? String == "object") + let properties = try #require(sent["properties"] as? [String: [String: Any]]) + #expect(properties["response"]?["type"] as? String == "array") + } else { + #expect((sent["type"] as? String)?.lowercased() == "array") + } + } + + /// A provider that doesn't apply the response schema can return content without the wrapper, + /// and an object with properties besides the wrapper property isn't unwrapped. + @Test(arguments: [false, true]) + func unwrappedResponseIsUsedAsSent(_ streaming: Bool) async throws { + SchemaURLProtocol.reset() + defer { SchemaURLProtocol.reset() } + let provider = Provider.chat + let operand = #"{"type":"number","value":1}"# + SchemaURLProtocol.enqueue( + json: try streaming ? provider.stream(text: operand) : json(provider.response(text: operand)) + ) + let session = provider.session() + let response = + try await streaming + ? session.streamResponse(to: "One", generating: Operand.self).collect() + : session.respond(to: "One", generating: Operand.self) + #expect(response.content == .number(1)) + + let extra = try GeneratedContent(json: #"{"response":{"type":"number","value":1},"note":"extra"}"#) + #expect(ProviderResponseSchema.unwrapping(extra, fromProperty: "response") == extra) + } + } + private func json(_ value: [String: Any]) throws -> String { String(decoding: try JSONSerialization.data(withJSONObject: value), as: UTF8.self) } From 57d10b684f1457d183dc974f033f8b95226dadbc Mon Sep 17 00:00:00 2001 From: mattt Date: Tue, 13 Oct 2026 05:02:42 -0700 Subject: [PATCH 7/9] Test an enum with associated values as a whole system model response The schema for such an enum has a union at the root, which the system model gets without a wrapper. --- .../SystemLanguageModelTests.swift | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/Tests/AnyLanguageModelTests/SystemLanguageModelTests.swift b/Tests/AnyLanguageModelTests/SystemLanguageModelTests.swift index b9135948..c18a7a77 100644 --- a/Tests/AnyLanguageModelTests/SystemLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/SystemLanguageModelTests.swift @@ -473,6 +473,21 @@ import Testing #expect(!response.content.steps.isEmpty) } + /// An enum with associated values can be the whole response, + /// with a union of its cases at the root of the schema. + @available(macOS 26.0, iOS 26.0, tvOS 26.0, visionOS 26.0, *) + @Test func guidedGenerationWithEnumPayloadRoot() async throws { + let session = LanguageModelSession(model: SystemLanguageModel.default) + + let response = try await session.respond( + to: "Add 2 to the counter.", + generating: CounterStep.self + ) + + let type = try response.rawContent.value(String.self, forProperty: "type") + #expect(["add", "reset", "report"].contains(type)) + } + @available(macOS 26.0, iOS 26.0, tvOS 26.0, visionOS 26.0, *) @Test func guidedGenerationWithInstructions() async throws { let session = LanguageModelSession( From 886a1f85062dfb7a20177fb778d47dc265e5adce Mon Sep 17 00:00:00 2001 From: mattt Date: Tue, 13 Oct 2026 05:09:18 -0700 Subject: [PATCH 8/9] Let local models choose the case of an enum with associated values Constrained generation for MLX, llama.cpp, and Core ML always generated the first choice of an anyOf. Now that such an enum's schema is a union of objects, local models could only produce its first case. When every choice is an object with a different constant for the same required property, like each case's type, the model now chooses that constant first and then generates that object's other properties. Other unions still use the first choice. --- .../AnyLanguageModel/GenerationSchema.swift | 3 +- .../Shared/ProviderResponseSchema.swift | 8 +-- .../Shared/StructuredGeneration.swift | 71 ++++++++++++++++++- .../StructuredGenerationTests.swift | 57 +++++++++++++++ 4 files changed, 128 insertions(+), 11 deletions(-) diff --git a/Sources/AnyLanguageModel/GenerationSchema.swift b/Sources/AnyLanguageModel/GenerationSchema.swift index 41a13d13..445ddb97 100644 --- a/Sources/AnyLanguageModel/GenerationSchema.swift +++ b/Sources/AnyLanguageModel/GenerationSchema.swift @@ -1155,7 +1155,8 @@ extension GenerationSchema { } } - private func resolving(_ node: Node, depth: Int) -> Node? { + /// Returns the node that a reference refers to, following any chain of references. + func resolving(_ node: Node, depth: Int = 0) -> Node? { guard depth < 64 else { return nil } guard case .ref(let name) = node else { return node } return defs[name].flatMap { resolving($0, depth: depth + 1) } diff --git a/Sources/AnyLanguageModel/Shared/ProviderResponseSchema.swift b/Sources/AnyLanguageModel/Shared/ProviderResponseSchema.swift index 30f5bf13..acd12113 100644 --- a/Sources/AnyLanguageModel/Shared/ProviderResponseSchema.swift +++ b/Sources/AnyLanguageModel/Shared/ProviderResponseSchema.swift @@ -73,13 +73,7 @@ struct ProviderResponseSchema { extension GenerationSchema { /// The root, after following references. var resolvedRoot: Node? { - var node = root - var seen: Set = [] - while case .ref(let name) = node { - guard seen.insert(name).inserted, let definition = defs[name] else { return nil } - node = definition - } - return node + resolving(root) } /// Returns the schema with its root as the one required property of an object. diff --git a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift index 5bb0880d..c4deb1c6 100644 --- a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift +++ b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift @@ -401,12 +401,77 @@ struct ConstrainedJSONGenerator { if variants.count == 1 { return try await generateNode(variants[0]) } + // The model chooses among objects that a constant property tells apart, + // like the cases of an enumeration with associated values. + if let union = discriminatedUnion(variants) { + return try await generateDiscriminatedUnion(union) + } // Choose the first variant to keep selection deterministic. return try await generateNode(variants[0]) } } + /// Objects that a required, constant string property tells apart. + private struct DiscriminatedUnion { + /// The name of the constant property. + var property: String + /// The value of the constant property in each object, with the object. + var cases: [(value: String, object: GenerationSchema.ObjectNode)] + } + + /// Returns the union of the variants if they're all objects + /// with a required property that has a different constant string value in each. + private func discriminatedUnion(_ variants: [GenerationSchema.Node]) -> DiscriminatedUnion? { + var objects: [GenerationSchema.ObjectNode] = [] + for variant in variants { + guard case .object(let object)? = schema.resolving(variant) else { return nil } + objects.append(object) + } + guard let first = objects.first else { return nil } + for property in first.required.sorted() { + let values = objects.compactMap { object -> String? in + guard object.required.contains(property), + case .string(let node)? = object.properties[property], + node.pattern == nil, + let choices = node.enumChoices, choices.count == 1 + else { return nil } + return choices[0] + } + if values.count == objects.count, Set(values).count == values.count { + return DiscriminatedUnion(property: property, cases: Array(zip(values, objects))) + } + } + return nil + } + + /// Generates one object of a discriminated union. + /// + /// The object starts with the constant property, + /// whose value the model chooses from the union's values, + /// and continues with the other properties of the object with that value. + private mutating func generateDiscriminatedUnion(_ union: DiscriminatedUnion) async throws -> String { + var output = try await emit("{") + output += try await emit("\"\(union.property)\":") + let value = try await generateString(GenerationSchema.StringNode(enumChoices: union.cases.map(\.value))) + output += value + guard let object = union.cases.first(where: { "\"\($0.value)\"" == value })?.object else { + throw ConstrainedGenerationError.tokenizationFailed + } + output += try await generateProperties(of: object, emitted: [union.property]) + return output + } + private mutating func generateObject(_ node: GenerationSchema.ObjectNode) async throws -> String { + let output = try await emit("{") + return output + (try await generateProperties(of: node, emitted: [])) + } + + /// Generates the properties of an object after its opening brace and the properties already emitted, + /// through its closing brace. + private mutating func generateProperties( + of node: GenerationSchema.ObjectNode, + emitted: Set + ) async throws -> String { // 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 @@ -416,10 +481,10 @@ struct ConstrainedJSONGenerator { // `}` 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) + var remainingKeys = Set(node.properties.keys).subtracting(emitted) let required = node.required - var output = try await emit("{") - var emittedAnyProperty = false + var output = "" + var emittedAnyProperty = !emitted.isEmpty while !remainingKeys.isEmpty { let missingRequired = required.intersection(remainingKeys) diff --git a/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift b/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift index 4d2e7758..5e54ad1c 100644 --- a/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift +++ b/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift @@ -860,6 +860,63 @@ struct StructuredGenerationTests { #expect(generator.generatedTokenCount == 2) } + // MARK: - Discriminated unions + + /// A union like the schema of an enumeration with associated values: + /// case `a` has the property `x`, and case `b` has the property `y`. + private func discriminatedUnionSchema() -> GenerationSchema { + func caseNode(_ type: String, property: String) -> GenerationSchema.Node { + .object( + GenerationSchema.ObjectNode( + description: nil, + properties: [ + "type": .string(.init(enumChoices: [type], isConstant: true)), + property: .string(.init(enumChoices: ["a"])), + ], + required: ["type", property] + ) + ) + } + return GenerationSchema( + root: .anyOf([.ref("A"), .ref("B")]), + defs: ["A": caseNode("a", property: "x"), "B": caseNode("b", property: "y")] + ) + } + + @Test(arguments: [("a", "x"), ("b", "y")]) + func modelChoosesTheObjectOfADiscriminatedUnion(_ type: String, _ property: String) async throws { + var maps = objectTokenMaps() + let quote = 0 + let comma = 1 + let colon = 4 + let aToken = 8 + let typeToken = 16 + let eosToken = 50 + maps.tokenToText[typeToken] = "type" + maps.textToTokens["\"type\":"] = [quote, typeToken, quote, colon] + let typeValue = try #require(maps.textToTokens[type]?.first) + let propertyToken = try #require(maps.textToTokens[property]?.first) + + // `{"type":` is forced; the model chooses the value of `type`, then the case's property. + let backend = MockTokenBackend( + tokenToText: maps.tokenToText, + textToTokens: maps.textToTokens, + eosToken: eosToken, + endTokens: [eosToken], + maximumTokens: 64, + samplingQueue: [ + typeValue, + comma, quote, propertyToken, quote, colon, + aToken, + ] + ) + + var generator = try ConstrainedJSONGenerator(backend: backend, schema: discriminatedUnionSchema()) + let result = try await generator.generate() + #expect(result == #"{"type":"\#(type)","\#(property)":"a"}"#) + #expect(generator.generatedTokenCount == backend.capture.tokens.count) + } + // MARK: - Decimal / number token mask private func numberTokenMaps() -> ( From 01a2bf4b2693784e23994e88333aa98ef2081ba6 Mon Sep 17 00:00:00 2001 From: mattt Date: Tue, 13 Oct 2026 05:47:36 -0700 Subject: [PATCH 9/9] Test that an enum case's object starts with its type Other unions follow the keys the model writes, so a payload key can come first. The object of an enumeration case can't start with anything but its type. --- .../StructuredGenerationTests.swift | 51 +++++++++++-------- 1 file changed, 30 insertions(+), 21 deletions(-) diff --git a/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift b/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift index 5e54ad1c..7a3bad63 100644 --- a/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift +++ b/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift @@ -883,38 +883,47 @@ struct StructuredGenerationTests { ) } - @Test(arguments: [("a", "x"), ("b", "y")]) - func modelChoosesTheObjectOfADiscriminatedUnion(_ type: String, _ property: String) async throws { + /// Generates the union of `discriminatedUnionSchema()`, sampling the tokens in `queue`. + private func generateDiscriminatedUnion( + sampling queue: (_ token: (String) -> Int) -> [Int] + ) async throws -> (result: String, generatedTokens: Int, decodedTokens: Int) { var maps = objectTokenMaps() - let quote = 0 - let comma = 1 - let colon = 4 - let aToken = 8 let typeToken = 16 - let eosToken = 50 maps.tokenToText[typeToken] = "type" - maps.textToTokens["\"type\":"] = [quote, typeToken, quote, colon] - let typeValue = try #require(maps.textToTokens[type]?.first) - let propertyToken = try #require(maps.textToTokens[property]?.first) - - // `{"type":` is forced; the model chooses the value of `type`, then the case's property. + maps.textToTokens["\"type\":"] = [0, typeToken, 0, 4] + maps.textToTokens[",\"type\":"] = [1, 0, typeToken, 0, 4] + let tokens = maps.textToTokens let backend = MockTokenBackend( tokenToText: maps.tokenToText, textToTokens: maps.textToTokens, - eosToken: eosToken, - endTokens: [eosToken], + eosToken: 50, + endTokens: [50], maximumTokens: 64, - samplingQueue: [ - typeValue, - comma, quote, propertyToken, quote, colon, - aToken, - ] + samplingQueue: queue { tokens[$0]![0] } ) - var generator = try ConstrainedJSONGenerator(backend: backend, schema: discriminatedUnionSchema()) let result = try await generator.generate() + return (result, generator.generatedTokenCount, backend.capture.tokens.count) + } + + @Test(arguments: [("a", "x"), ("b", "y")]) + func modelChoosesTheObjectOfADiscriminatedUnion(_ type: String, _ property: String) async throws { + // `{"type":` is forced; the model chooses the value of `type`, then the case's property. + let (result, generated, decoded) = try await generateDiscriminatedUnion { token in + [token(type), token(","), token("\""), token(property), token("\""), token(":"), token("a")] + } #expect(result == #"{"type":"\#(type)","\#(property)":"a"}"#) - #expect(generator.generatedTokenCount == backend.capture.tokens.count) + #expect(generated == decoded) + } + + /// Unlike other unions, whose keys come in the order the model writes them, + /// the object of an enumeration case starts with its `type`. + @Test func caseTypeComesBeforeTheOtherKeys() async throws { + await #expect(throws: ConstrainedGenerationError.self) { + try await generateDiscriminatedUnion { token in + [token("\""), token("y"), token("\""), token(":"), token("a")] + } + } } // MARK: - Decimal / number token mask