Port JSONCompleter and drop the PartialJSONDecoder dependency

Package.swift

@@ -41,7 +41,6 @@ let package = Package(
         ),
         .package(url: "https://github.com/mattt/JSONSchema", from: "1.3.0"),
         .package(url: "https://github.com/mattt/llama.swift", .upToNextMajor(from: "2.10549.0")),
-        .package(url: "https://github.com/mattt/PartialJSONDecoder", from: "1.0.0"),
         .package(url: "https://github.com/ml-explore/mlx-swift-lm", from: "3.31.4"),
         .package(url: "https://github.com/swiftlang/swift-syntax", from: "602.0.0"),
         .package(url: "https://github.com/swift-server/async-http-client.git", from: "1.24.0"),
@@ -53,7 +52,6 @@ let package = Package(
                 .target(name: "AnyLanguageModelMacros"),
                 .product(name: "EventSource", package: "EventSource"),
                 .product(name: "JSONSchema", package: "JSONSchema"),
-                .product(name: "PartialJSONDecoder", package: "PartialJSONDecoder"),
                 .product(
                     name: "MLXLLM",
                     package: "mlx-swift-lm",

Sources/AnyLanguageModel/GeneratedContent.swift

@@ -179,33 +179,17 @@ public struct GeneratedContent: Sendable, Equatable, Generable, CustomDebugStrin
             return
         }
 
-        // Handle incomplete JSON by attempting to complete it
-        let completedJSON = String(decoding: data, as: UTF8.self)
-            .trimmingCharacters(in: .whitespacesAndNewlines)
-
-        // Try adding closing braces/brackets to make it valid
-        var attempts: [String] = [completedJSON]
-
-        // If it looks like an incomplete object, try closing it
-        if completedJSON.hasPrefix("{") && !completedJSON.hasSuffix("}") {
-            attempts.append(completedJSON + "}")
-            attempts.append(completedJSON + "\"\"}")  // incomplete string value
-        }
-
-        // If it looks like an incomplete array, try closing it
-        if completedJSON.hasPrefix("[") && !completedJSON.hasSuffix("]") {
-            attempts.append(completedJSON + "]")
-        }
-
-        for attempt in attempts {
-            if let parsed = try? JSONSerialization.jsonObject(with: Data(attempt.utf8), options: [.fragmentsAllowed]) {
-                self = try Self.fromJSONObject(parsed)
-                return
-            }
+        // Handle incomplete JSON by completing it and parsing again
+        let json = String(decoding: data, as: UTF8.self)
+        if let completed = try? JSONCompleter().complete(json),
+            let parsed = try? JSONSerialization.jsonObject(with: Data(completed.utf8), options: [.fragmentsAllowed])
+        {
+            self = try Self.fromJSONObject(parsed)
+            return
         }
 
         // If all else fails, treat it as a string
-        self.init(kind: .string(completedJSON))
+        self.init(kind: .string(json.trimmingCharacters(in: .whitespacesAndNewlines)))
     }
 
     private static func fromJSONObject(_ value: Any) throws -> GeneratedContent {

Sources/AnyLanguageModel/Models/SystemLanguageModel.swift

@@ -1,7 +1,6 @@
 #if canImport(FoundationModels)
     import FoundationModels
     import Foundation
-    import PartialJSONDecoder
 
     import JSONSchema
 
@@ -788,9 +787,7 @@
                 return finalize(content: content)
             } catch {
                 // Attempt partial JSON decoding before surfacing an error.
-                let decoder = PartialJSONDecoder()
-                let jsonString = fmResponse.content.jsonString
-                if let partialContent = try? decoder.decode(GeneratedContent.self, from: jsonString).value,
+                if let partialContent = try? GeneratedContent(json: fmResponse.content.jsonString),
                     let content = try? type.init(partialContent)
                 {
                     return finalize(content: content)
@@ -870,7 +867,6 @@
 
                 func processStructuredStream(_ fmSession: FoundationModels.LanguageModelSession) async {
                     let fmSchema = FoundationModels.GenerationSchema(schema)
-                    let partialDecoder = PartialJSONDecoder()
                     let fmStream = fmSession.streamResponse(
                         to: fmPrompt,
                         schema: fmSchema,
@@ -898,12 +894,7 @@
                                     lastLength: &lastLength
                                 )
 
-                                let jsonString = accumulatedText
-                                if let partialContent = try? partialDecoder.decode(
-                                    GeneratedContent.self,
-                                    from: jsonString
-                                )
-                                .value {
+                                if let partialContent = try? GeneratedContent(json: accumulatedText) {
                                     let partial: Content.PartiallyGenerated? = try? .init(partialContent)
                                     if let partial {
                                         continuation.yield(.init(content: partial, rawContent: partialContent))
@@ -937,11 +928,7 @@
                                 ?? GeneratedContent(jsonString)
 
                             // Prefer partial decoding so we can surface intermediate snapshots.
-                            if let partialContent = try? partialDecoder.decode(
-                                GeneratedContent.self,
-                                from: jsonString
-                            )
-                            .value {
+                            if let partialContent = try? GeneratedContent(json: jsonString) {
                                 let partial: Content.PartiallyGenerated? = try? .init(partialContent)
                                 if let partial {
                                     continuation.yield(.init(content: partial, rawContent: partialContent))