From 3c7a0bdffcd3edcafa4bbe673b4b5a53770cbaff Mon Sep 17 00:00:00 2001 From: Jaleel Akbashev Date: Tue, 6 Oct 2026 10:50:07 +0200 Subject: [PATCH] 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(