Generate enums with associated values as Foundation Models does

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<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)
     }
 

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<String> = [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) }

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,

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<Content>()
+                    var state = StreamingResponseState<Content>(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

Sources/AnyLanguageModel/Models/OpenAILanguageModel.swift

@@ -492,7 +492,8 @@ public struct OpenAILanguageModel: LanguageModel {
         options: GenerationOptions,
         session: LanguageModelSession
     ) async throws -> LanguageModelSession.Response<Content> 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<Content> 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<Content> where Content: Generable {
+        // Structured outputs need an object root.
+        let responseSchema = ProviderResponseSchema(schema, generating: type, wrapping: .nonObjectRoots)
         let stream = AsyncThrowingStream<LanguageModelSession.ResponseStream<Content>.Snapshot, any Error> {
             continuation in
             let task = Task {
                 do {
                     var inFlightMessages: [OpenAIMessage] = []
-                    var state = StreamingResponseState<Content>()
+                    var state = StreamingResponseState<Content>(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

Sources/AnyLanguageModel/Models/OpenResponsesLanguageModel.swift

@@ -486,12 +486,14 @@ public struct OpenResponsesLanguageModel: LanguageModel {
         options: GenerationOptions
     ) -> sending LanguageModelSession.ResponseStream<Content> 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<LanguageModelSession.ResponseStream<Content>.Snapshot, any Error> {
             continuation in
             let task = Task {
                 do {
                     var inFlightMessages: [OpenResponsesMessage] = []
-                    var state = StreamingResponseState<Content>()
+                    var state = StreamingResponseState<Content>(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<Content> 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,