Initial implementation for structured chat messages

Author: ibrahimcetin

Libraries/MLXLLM/LLMModelFactory.swift

@@ -199,23 +199,37 @@ private struct LLMUserInputProcessor: UserInputProcessor {
 
     let tokenizer: Tokenizer
     let configuration: ModelConfiguration
+    let messageGenerator: MessageGenerator
 
-    internal init(tokenizer: any Tokenizer, configuration: ModelConfiguration) {
+    internal init(
+        tokenizer: any Tokenizer, configuration: ModelConfiguration,
+        messageGenerator: MessageGenerator
+    ) {
         self.tokenizer = tokenizer
         self.configuration = configuration
+        self.messageGenerator = messageGenerator
     }
 
     func prepare(input: UserInput) throws -> LMInput {
+        let messages =
+            switch input.prompt {
+            case .text(let text):
+                messageGenerator.generate(messages: [.user(text)])
+            case .messages(let messages):
+                messages
+            case .chat(let messages):
+                messageGenerator.generate(messages: messages)
+            }
+
         do {
-            let messages = input.prompt.asMessages()
             let promptTokens = try tokenizer.applyChatTemplate(
                 messages: messages, tools: input.tools, additionalContext: input.additionalContext)
             return LMInput(tokens: MLXArray(promptTokens))
         } catch {
             // #150 -- it might be a TokenizerError.chatTemplate("No chat template was specified")
             // but that is not public so just fall back to text
-            let prompt = input.prompt
-                .asMessages()
+            let prompt =
+                messages
                 .compactMap { $0["content"] as? String }
                 .joined(separator: ". ")
             let promptTokens = tokenizer.encode(text: prompt)
@@ -273,7 +287,9 @@ public class LLMModelFactory: ModelFactory {
 
         return .init(
             configuration: configuration, model: model,
-            processor: LLMUserInputProcessor(tokenizer: tokenizer, configuration: configuration),
+            processor: LLMUserInputProcessor(
+                tokenizer: tokenizer, configuration: configuration,
+                messageGenerator: DefaultMessageGenerator()),
             tokenizer: tokenizer)
     }
 

Libraries/MLXLMCommon/Chat.swift

@@ -0,0 +1,91 @@
+// Copyright © 2025 Apple Inc.
+
+public enum Chat {
+    public struct Message {
+        /// The role of the message sender.
+        public let role: Role
+
+        /// The content of the message.
+        public let content: String
+
+        /// Array of image data associated with the message.
+        public let images: [UserInput.Image]
+
+        /// Array of video data associated with the message.
+        public let videos: [UserInput.Video]
+
+        public static func system(
+            _ content: String, images: [UserInput.Image] = [], videos: [UserInput.Video] = []
+        ) -> Self {
+            Self(role: .system, content: content, images: images, videos: videos)
+        }
+
+        public static func assistant(
+            _ content: String, images: [UserInput.Image] = [], videos: [UserInput.Video] = []
+        ) -> Self {
+            Self(role: .assistant, content: content, images: images, videos: videos)
+        }
+
+        public static func user(
+            _ content: String, images: [UserInput.Image] = [], videos: [UserInput.Video] = []
+        ) -> Self {
+            Self(role: .user, content: content, images: images, videos: videos)
+        }
+
+        public enum Role: String {
+            case user
+            case assistant
+            case system
+        }
+    }
+}
+
+public protocol MessageGenerator {
+    /// Returns [String: Any] aka Message
+    func generate(message: Chat.Message) -> Message
+}
+
+extension MessageGenerator {
+    /// Returns array of [String: Any] aka Message
+    public func generate(messages: [Chat.Message]) -> [Message] {
+        var rawMessages: [Message] = []
+
+        for message in messages {
+            let raw = generate(message: message)
+            rawMessages.append(raw)
+        }
+
+        return rawMessages
+    }
+}
+
+public struct DefaultMessageGenerator: MessageGenerator {
+    public init() {}
+
+    public func generate(message: Chat.Message) -> Message {
+        [
+            "role": message.role.rawValue,
+            "content": message.content,
+        ]
+    }
+}
+
+public struct Qwen2VLMessageGenerator: MessageGenerator {
+    public init() {}
+
+    public func generate(message: Chat.Message) -> Message {
+        [
+            "role": message.role.rawValue,
+            "content": [
+                ["type": "text", "text": message.content]
+            ]
+                // Messages format for Qwen 2 VL, Qwen 2.5 VL. May need to be adapted for other models.
+                + message.images.map { _ in
+                    ["type": "image"]
+                }
+                + message.videos.map { _ in
+                    ["type": "video"]
+                },
+        ]
+    }
+}

Libraries/MLXLMCommon/UserInput.swift

@@ -17,22 +17,16 @@ public struct UserInput: Sendable {
     public enum Prompt: Sendable, CustomStringConvertible {
         case text(String)
         case messages([Message])
-
-        public func asMessages() -> [Message] {
-            switch self {
-            case .text(let text):
-                return [["role": "user", "content": text]]
-            case .messages(let messages):
-                return messages
-            }
-        }
+        case chat([Chat.Message])
 
         public var description: String {
             switch self {
             case .text(let text):
                 return text
             case .messages(let messages):
                 return messages.map { $0.description }.joined(separator: "\n")
+            case .chat(let messages):
+                return messages.map(\.content).joined(separator: "\n")
             }
         }
     }
@@ -156,6 +150,18 @@ public struct UserInput: Sendable {
         self.additionalContext = additionalContext
     }
 
+    public init(
+        messages: [Chat.Message], images: [Image] = [Image](), videos: [Video] = [Video](),
+        tools: [ToolSpec]? = nil,
+        additionalContext: [String: Any]? = nil
+    ) {
+        self.prompt = .chat(messages)
+        self.images = images
+        self.videos = videos
+        self.tools = tools
+        self.additionalContext = additionalContext
+    }
+
     public init(
         prompt: Prompt, images: [Image] = [Image](), processing: Processing = .init(),
         tools: [ToolSpec]? = nil, additionalContext: [String: Any]? = nil

Libraries/MLXVLM/Models/Idefics3.swift

@@ -817,10 +817,19 @@ public class Idefics3Processor: UserInputProcessor {
         self.tokenizer = tokenizer
     }
 
-    public func prepare(input: UserInput) throws -> LMInput {
-
-        let prompt = input.prompt.asMessages().last?["content"] as? String ?? ""
+    private func prompt(from userInput: UserInput) -> String {
+        switch userInput.prompt {
+        case .text(let text):
+            text
+        case .messages(let messages):
+            messages.last?["content"] as? String ?? ""
+        case .chat(let messages):
+            messages.last?.content ?? ""
+        }
+    }
 
+    public func prepare(input: UserInput) throws -> LMInput {
+        let prompt = prompt(from: input)
         if input.images.isEmpty {
             // No image scenario
             let tokens = try tokenizer.encode(text: prompt)

Libraries/MLXVLM/Models/Paligemma.swift

@@ -478,7 +478,7 @@ public class PaligGemmaProcessor: UserInputProcessor {
         }
 
         // this doesn't have a chat template so just use the last message.
-        var prompt = input.prompt.asMessages().last?["content"] as? String ?? ""
+        var prompt = prompt(from: input)
 
         // based on transformers/processing_paligemma
         let count = input.images.count * config.imageSequenceLength
@@ -495,6 +495,17 @@ public class PaligGemmaProcessor: UserInputProcessor {
         return LMInput(text: .init(tokens: promptArray, mask: mask), image: .init(pixels: pixels))
     }
 
+    private func prompt(from userInput: UserInput) -> String {
+        switch userInput.prompt {
+        case .text(let text):
+            text
+        case .messages(let messages):
+            messages.last?["content"] as? String ?? ""
+        case .chat(let messages):
+            messages.last?.content ?? ""
+        }
+    }
+
 }
 
 // MARK: - Model

Libraries/MLXVLM/Models/Qwen2VL.swift

@@ -696,7 +696,16 @@ public class Qwen2VLProcessor: UserInputProcessor {
     }
 
     public func prepare(input: UserInput) async throws -> LMInput {
-        let messages = input.prompt.asMessages()
+        let generator = Qwen2VLMessageGenerator()
+        let messages =
+            switch input.prompt {
+            case .text(let text):
+                generator.generate(messages: [.user(text)])
+            case .messages(let messages):
+                messages
+            case .chat(let messages):
+                generator.generate(messages: messages)
+            }
         var promptTokens = try tokenizer.applyChatTemplate(messages: messages)
 
         // Text-only input