Implement prewarm for MLXLanguageModel

Author: noorbhatia

Sources/AnyLanguageModel/Models/MLXLanguageModel.swift

@@ -357,6 +357,51 @@ import Foundation
 
             return LanguageModelSession.ResponseStream(stream: stream)
         }
+
+        /// Prewarms the model for the given session and optional prompt prefix.
+        public func prewarm(
+            for session: LanguageModelSession,
+            promptPrefix: Prompt?
+        ) {
+            let modelId = self.modelId
+            let hub = self.hub
+            let directory = self.directory
+
+            let instructions = session.instructions?.description
+            let tools = session.tools
+
+            Task {
+
+                let context = try await loadContext(modelId: modelId, hub: hub, directory: directory)
+
+                // Build chat history similar to respond() to prime the cache effectively
+                var chat: [MLXLMCommon.Chat.Message] = []
+
+                // Add system instructions if present
+                if let instructions, !instructions.isEmpty {
+                    chat.append(.init(role: .system, content: instructions))
+                }
+
+                // Add prompt prefix or minimal user message
+                let promptText = promptPrefix?.description ?? "."
+                chat.append(.init(role: .user, content: promptText))
+
+                // Convert tools to MLX format
+                let toolSpecs: [ToolSpec]? =
+                    tools.isEmpty
+                    ? nil
+                    : tools.map { convertToolToMLXSpec($0) }
+
+                let userInput = MLXLMCommon.UserInput(
+                    chat: chat,
+                    processing: .init(resize: .init(width: 512, height: 512)),
+                    tools: toolSpecs
+                )
+
+                // Prepare input - triggers tokenization and processor initialization
+                _ = try await context.processor.prepare(input: userInput)
+            }
+        }
     }
 
     // MARK: - Options Mapping