Incorporate feedback from review

Author: mattt

README.md

@@ -502,6 +502,12 @@ mlxOptions.kvCache = .init(
 )
 // Apply a deterministic preprocessing step for image inputs.
 mlxOptions.userInputProcessing = .resize(to: CGSize(width: 512, height: 512))
+// Inject extra template context consumed by model-specific chat templates.
+mlxOptions.additionalContext = [
+    "user_name": .string("Alice"),
+    "turn_count": .int(3),
+    "verbose": .bool(true),
+]
 options[custom: MLXLanguageModel.self] = mlxOptions
 
 let response = try await session.respond(
@@ -511,11 +517,14 @@ let response = try await session.respond(
 ```
 
 You can specify `userInputProcessing` to enforce a consistent image
-preprocessing step 
+preprocessing step
 (for example, fixed dimensions for predictable latency, memory usage, and vision behavior).
 By default, images are passed through without an explicit resize override
 (`resize: nil`), so MLX applies its default media processing behavior.
 
+You can also set `additionalContext` to provide extra JSON template variables
+for model-specific chat templates.
+
 GPU cache behavior can be configured when creating the model:
 
 ```swift

Sources/AnyLanguageModel/Models/MLXLanguageModel.swift

@@ -281,7 +281,7 @@ import Foundation
             }
 
             /// Additional key-value pairs injected into the chat template rendering context.
-            public var additionalContext: [String: MLXLMCommon.JSONValue]?
+            public var additionalContext: [String: AnyLanguageModel.JSONValue]?
 
             var additionalContextForUserInput: [String: any Sendable]? {
                 additionalContext?.mapValues { $0.toSendable() }
@@ -298,7 +298,7 @@ import Foundation
             public init(
                 kvCache: KVCache,
                 userInputProcessing: UserInputProcessing?,
-                additionalContext: [String: MLXLMCommon.JSONValue]?
+                additionalContext: [String: AnyLanguageModel.JSONValue]?
             ) {
                 self.kvCache = kvCache
                 self.additionalContext = additionalContext
@@ -1882,7 +1882,7 @@ import Foundation
             return sampledToken.item(Int.self)
         }
     }
-    extension MLXLMCommon.JSONValue {
+    extension AnyLanguageModel.JSONValue {
         /// Recursively converts a `JSONValue` to its primitive Swift equivalent.
         func toSendable() -> any Sendable {
             switch self {

Sources/AnyLanguageModel/Shared/JSONValue.swift

@@ -0,0 +1,4 @@
+import enum JSONSchema.JSONValue
+
+/// A type-safe representation of JSON values used by AnyLanguageModel APIs.
+public typealias JSONValue = JSONSchema.JSONValue

Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift

@@ -4,7 +4,6 @@ import Testing
 @testable import AnyLanguageModel
 
 #if MLX
-    import MLXLMCommon
     private let shouldRunMLXTests = {
         // Enable when explicitly requested via environment variable
         if ProcessInfo.processInfo.environment["ENABLE_MLX_TESTS"] != nil {
@@ -273,9 +272,9 @@ import Testing
             )
             var custom = MLXLanguageModel.CustomGenerationOptions.default
             custom.additionalContext = [
-                "user_name": MLXLMCommon.JSONValue.string("Alice"),
-                "turn_count": MLXLMCommon.JSONValue.int(3),
-                "verbose": MLXLMCommon.JSONValue.bool(true),
+                "user_name": JSONValue.string("Alice"),
+                "turn_count": JSONValue.int(3),
+                "verbose": JSONValue.bool(true),
             ]
             options[custom: MLXLanguageModel.self] = custom