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/Sources/AnyLanguageModel/GenerationSchema.swift b/Sources/AnyLanguageModel/GenerationSchema.swift index 165cf3b1..a5a0c76a 100644 --- a/Sources/AnyLanguageModel/GenerationSchema.swift +++ b/Sources/AnyLanguageModel/GenerationSchema.swift @@ -431,6 +431,79 @@ public struct GenerationSchema: Equatable, Codable, CustomDebugStringConvertible self.defs = [:] } + /// The property that names the case in the object of an enumeration case, + /// as in Foundation Models. + static let caseTypeProperty = "type" + + /// 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] = [:] + + // 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 { + // Like Foundation Models, `type` is a constant. + // Providers whose schemas don't support `const` get a one-choice `enum`. + let typeProperty = Self.caseTypeProperty + var props: [String: Node] = [typeProperty: .string(StringNode(enumChoices: [name], isConstant: true))] + var required: Set = [typeProperty] + for property in properties { + props[property.name] = property.node + if !property.isOptional { + required.insert(property.name) + } + } + // 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, + properties: props, + required: required, + propertyOrder: [typeProperty] + 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: @@ -535,7 +608,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 } @@ -1087,7 +1160,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/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index 552c3fe0..61ae70c9 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/Sources/AnyLanguageModel/LanguageModelSession.swift @@ -871,7 +871,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, @@ -891,7 +891,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, @@ -908,7 +908,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, @@ -925,7 +925,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, @@ -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( to prompt: String, schema: GenerationSchema, @@ -980,7 +980,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..acd12113 --- /dev/null +++ b/Sources/AnyLanguageModel/Shared/ProviderResponseSchema.swift @@ -0,0 +1,89 @@ +/// 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? { + resolving(root) + } + + /// 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/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift index d8049eb5..ccf66f40 100644 --- a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift +++ b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift @@ -502,6 +502,8 @@ struct ConstrainedJSONGenerator { /// /// With several variants, the mask offers the keys of every variant still in play, and /// each emitted key keeps only the variants that declare it. + /// When the variants are the cases of an enumeration with associated values, + /// the object starts with the case's `type`, whose value chooses the case. /// /// - Returns: The object, and the positions in `variants` of the variants it still matches. private mutating func generateObject( @@ -521,6 +523,14 @@ struct ConstrainedJSONGenerator { var emittedKeys = Set() var output = try await emit("{") + // Like Foundation Models, the object of an enumeration case starts with its `type`, + // whose value chooses the case. Other unions follow the keys the model writes. + if variants.count > 1, let key = try caseTypeKey(of: variants) { + output += try await emit("\"\(key)\":") + emittedKeys.insert(key) + output += try await generateValue(forKey: key, of: variants, narrowing: &live) + } + while true { let liveVariants = live.map { variants[$0] } let remainingKeys = liveVariants.reduce(into: Set()) { $0.formUnion($1.properties.keys) } @@ -730,6 +740,24 @@ struct ConstrainedJSONGenerator { } } + /// The key that names the case, if every variant is the object of a different case + /// of an enumeration with associated values. + /// + /// That's a required `type` with a different single string choice in each variant, + /// as `GenerationSchema(type:description:cases:)` writes it. + private func caseTypeKey(of variants: [GenerationSchema.ObjectNode]) throws -> String? { + let key = GenerationSchema.caseTypeProperty + var names = Set() + for variant in variants { + guard variant.required.contains(key), + let node = variant.properties[key], + let choices = try stringChoices(of: node), choices.count == 1, + names.insert(choices[0]).inserted + else { return nil } + } + return key + } + /// The choices of a string enum without a pattern, through references. private func stringChoices(of node: GenerationSchema.Node) throws -> [String]? { let flattened = try flattenedVariants([node]) diff --git a/Sources/AnyLanguageModelMacros/GenerableMacro.swift b/Sources/AnyLanguageModelMacros/GenerableMacro.swift index a7c74697..8d219b24 100644 --- a/Sources/AnyLanguageModelMacros/GenerableMacro.swift +++ b/Sources/AnyLanguageModelMacros/GenerableMacro.swift @@ -1089,68 +1089,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: ", "))]") + ) } } """ @@ -1184,116 +1147,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?, @@ -1303,28 +1156,21 @@ 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)") - ]) - """ - } else { - return generateMultipleValueSerialization( - caseName: enumCase.name, - associatedValues: enumCase.associatedValues - ) - } + let payload = enumCase.payload + let bindings = payload.indices.map { "let value\($0)" }.joined(separator: ", ") + let pattern = payload.isEmpty ? ".\(enumCase.name)" : ".\(enumCase.name)(\(bindings))" + // 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 + "properties[\"\(value.property)\"] = value\(index).generatedContent" + }.joined(separator: "\n ") + let keys = (["type"] + payload.map(\.property)).map { "\"\($0)\"" }.joined(separator: ", ") + return """ + case \(pattern): + \(payload.isEmpty ? "let" : "var") properties: [String: GeneratedContent] = ["type": GeneratedContent("\(enumCase.name)")] + \(assignments) + return GeneratedContent(kind: .structure(properties: properties, orderedKeys: [\(keys)])) + """ }.joined(separator: "\n ") return DeclSyntax( @@ -1353,61 +1199,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: """ @@ -1426,23 +1217,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: """ @@ -1450,9 +1230,8 @@ public struct GenerableMacro: MemberMacro, ExtensionMacro { return GenerationSchema( type: Self.self, description: \(description.map { "\"\($0)\"" } ?? "\"Generated \(enumName)\""), - properties: [ - \(caseProperty), - \(valueProperty) + cases: [ + \(caseSchemas) ] ) } @@ -1509,13 +1288,17 @@ 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. + /// + /// 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 + (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..5d4813d2 --- /dev/null +++ b/Tests/AnyLanguageModelTests/GenerableEnumPayloadFoundationModelsTests.swift @@ -0,0 +1,115 @@ +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)) + } + } + + /// 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 { + 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 new file mode 100644 index 00000000..6c453068 --- /dev/null +++ b/Tests/AnyLanguageModelTests/GenerableEnumPayloadTests.swift @@ -0,0 +1,236 @@ +import Foundation +import Testing + +@testable import AnyLanguageModel + +// `GenerableEnumPayloadFoundationModelsTests` compares these enums +// with the same enums in Foundation Models. + +@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? + +/// 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) + 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. +@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 aMissingOptionalValueIsNil() throws { + #expect(try Step(GeneratedContent(json: #"{"type": "note"}"#)) == .note(nil)) + #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)) + } + + /// 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"}"#)) + } + } + + @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"]?["const"] 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"]) + } + + /// 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 + ) 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"]?["const"] 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) + } +} 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", 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) } 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 diff --git a/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift b/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift index 4d2e7758..7a3bad63 100644 --- a/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift +++ b/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift @@ -860,6 +860,72 @@ 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")] + ) + } + + /// 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 typeToken = 16 + maps.tokenToText[typeToken] = "type" + 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: 50, + endTokens: [50], + maximumTokens: 64, + 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(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 private func numberTokenMaps() -> ( diff --git a/Tests/AnyLanguageModelTests/SystemLanguageModelTests.swift b/Tests/AnyLanguageModelTests/SystemLanguageModelTests.swift index 44690e50..c18a7a77 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,34 @@ 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) + } + + /// 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(