Fix Gemma 4 system message and modality order (#211)

* Align Gemma4 messages with chat template

* Respect Gemma 4 modality order

Author: adrgrondin

Libraries/MLXVLM/Models/Gemma4.swift

@@ -1769,6 +1769,32 @@ public final class Gemma4: Module, VLMModel, KVCacheDimensionProvider {
 
 // MARK: - Processor
 
+public struct Gemma4MessageGenerator: MessageGenerator {
+    public init() {}
+
+    public func generate(message: Chat.Message) -> MLXLMCommon.Message {
+        if message.role == .system {
+            [
+                "role": message.role.rawValue,
+                "content": message.content,
+            ]
+        } else {
+            [
+                "role": message.role.rawValue,
+                "content": message.images.map { _ in
+                    ["type": "image"]
+                }
+                    + message.videos.map { _ in
+                        ["type": "video"]
+                    }
+                    + [
+                        ["type": "text", "text": message.content]
+                    ],
+            ]
+        }
+    }
+}
+
 public struct Gemma4Processor: UserInputProcessor {
     private let config: Gemma4ProcessorConfiguration
     private let tokenizer: any Tokenizer
@@ -1805,7 +1831,7 @@ public struct Gemma4Processor: UserInputProcessor {
     }
 
     public func prepare(input: UserInput) async throws -> LMInput {
-        let messages = Qwen2VLMessageGenerator().generate(from: input)
+        let messages = Gemma4MessageGenerator().generate(from: input)
 
         var promptTokens = try tokenizer.applyChatTemplate(
             messages: messages, tools: input.tools,

Tests/MLXLMTests/UserInputTests.swift

@@ -95,6 +95,33 @@ public class UserInputTests: XCTestCase {
         assertEqual(expected, messages)
     }
 
+    public func testGemma4ConversionText() {
+        let chat: [Chat.Message] = [
+            .system("You are a useful agent."),
+            .user("Tell me a story."),
+        ]
+
+        let messages = Gemma4MessageGenerator().generate(messages: chat)
+
+        let expected: [[String: any Sendable]] = [
+            [
+                "role": "system",
+                "content": "You are a useful agent.",
+            ],
+            [
+                "role": "user",
+                "content": [
+                    [
+                        "type": "text",
+                        "text": "Tell me a story.",
+                    ]
+                ],
+            ],
+        ]
+
+        assertEqual(expected, messages)
+    }
+
     // MARK: - Mistral3 Message Generator Tests
 
     public func testMistral3ConversionText() {
@@ -230,4 +257,41 @@ public class UserInputTests: XCTestCase {
         XCTAssertEqual(userInput.images.count, 1)
     }
 
+    public func testGemma4ConversionImage() {
+        let chat: [Chat.Message] = [
+            .system("You are a useful agent."),
+            .user(
+                "What is this?",
+                images: [
+                    .url(
+                        URL(
+                            string: "https://opensource.apple.com/images/projects/mlx.f5c59d8b.png")!
+                    )
+                ]),
+        ]
+
+        let messages = Gemma4MessageGenerator().generate(messages: chat)
+
+        let expected: [[String: any Sendable]] = [
+            [
+                "role": "system",
+                "content": "You are a useful agent.",
+            ],
+            [
+                "role": "user",
+                "content": [
+                    [
+                        "type": "image"
+                    ],
+                    [
+                        "type": "text",
+                        "text": "What is this?",
+                    ],
+                ],
+            ],
+        ]
+
+        assertEqual(expected, messages)
+    }
+
 }