From a3cb31cdd32c6166443661d9d6c4e7dbe8cb0b58 Mon Sep 17 00:00:00 2001 From: Anthony Drendel Date: Tue, 7 Jul 2026 23:38:44 +0200 Subject: [PATCH] Enable Gemma 4 unified thinking by default - Update mlx-swift-lm to fix tool schemas for Gemma 4 --- Package.resolved | 5 ++--- Package.swift | 2 +- Sources/SHLLM/LLM.swift | 20 ++++++++----------- Tests/SHLLMTests/Models/Gemma4-12BTests.swift | 8 +------- Tests/SHLLMTests/Models/Gemma4-E2BTests.swift | 8 +------- Tests/SHLLMTests/Models/Gemma4-E4BTests.swift | 8 +------- 6 files changed, 14 insertions(+), 37 deletions(-) diff --git a/Package.resolved b/Package.resolved index e535898..90f757a 100644 --- a/Package.resolved +++ b/Package.resolved @@ -1,5 +1,5 @@ { - "originHash" : "249e833911a08229425856e4e4bf99018e562b937b235fb4ed690e5a9f732467", + "originHash" : "0448e12695cfc33f4ac3802d6174cf2f1caab9df0693adda9549f2fc494926a6", "pins" : [ { "identity" : "mlx-swift", @@ -15,8 +15,7 @@ "kind" : "remoteSourceControl", "location" : "https://github.com/shareup/mlx-swift-lm", "state" : { - "revision" : "fa8400f8b369df8fdd845d4958339788cca83520", - "version" : "0.0.16" + "revision" : "d414b738a56f6a0f92800f0ff1f663167bb0f871" } }, { diff --git a/Package.swift b/Package.swift index df2612a..49e1f6f 100644 --- a/Package.swift +++ b/Package.swift @@ -17,7 +17,7 @@ let package = Package( ), .package( url: "https://github.com/shareup/mlx-swift-lm", - from: "0.0.16" + revision: "d414b738a56f6a0f92800f0ff1f663167bb0f871" ), .package( url: "https://github.com/DePasqualeOrg/swift-tokenizers", diff --git a/Sources/SHLLM/LLM.swift b/Sources/SHLLM/LLM.swift index af2cae9..d5ce77e 100644 --- a/Sources/SHLLM/LLM.swift +++ b/Sources/SHLLM/LLM.swift @@ -633,12 +633,7 @@ extension LLM where Model == Gemma4 { /// **gemma-4-e2b** and **gemma-4-e4b** /// /// To enable thinking, set - /// `UserInput(additionalContext: ["enable_thinking": true])` - /// and add `<|think|>` to the **beginning** and **end** of - /// the system prompt. (The documentation says to just add it to - /// the beginning of the system prompt, but testing has shown it - /// to be more reliable when added to the beginning and end of the - /// system prompt. + /// `UserInput(additionalContext: ["enable_thinking": true])`. /// /// https://huggingface.co/google/gemma-4-12B-it#2-thinking-mode-configuration public static func gemma4( @@ -701,12 +696,7 @@ extension LLM where Model == Gemma4Unified { /// **gemma-4-12b** /// /// To enable thinking, set - /// `UserInput(additionalContext: ["enable_thinking": true])` - /// and add `<|think|>` to the **beginning** and **end** of - /// the system prompt. (The documentation says to just add it to - /// the beginning of the system prompt, but testing has shown it - /// to be more reliable when added to the beginning and end of the - /// system prompt. + /// `UserInput(additionalContext: ["enable_thinking": true])`. /// /// https://huggingface.co/google/gemma-4-12B-it#2-thinking-mode-configuration public static func gemma4Unified( @@ -717,6 +707,12 @@ extension LLM where Model == Gemma4Unified { maxOutputTokenCount: Int? = nil ) throws -> LLM { try SHLLM.assertSupportedDevice + var input = input + var additionalContext = input.additionalContext ?? [:] + if additionalContext["enable_thinking"] == nil { + additionalContext["enable_thinking"] = true + } + input.additionalContext = additionalContext return .init( directory: directory, input: input, diff --git a/Tests/SHLLMTests/Models/Gemma4-12BTests.swift b/Tests/SHLLMTests/Models/Gemma4-12BTests.swift index c90011f..63f0fec 100644 --- a/Tests/SHLLMTests/Models/Gemma4-12BTests.swift +++ b/Tests/SHLLMTests/Models/Gemma4-12BTests.swift @@ -104,7 +104,7 @@ struct Gemma4_12BTests { let input = UserInput( chat: [ .system( - "You are a weather assistant who must use the get_current_weather tool to fetch weather data for any location the user asks about.<|think|>" + "You are a weather assistant who must use the get_current_weather tool to fetch weather data for any location the user asks about." ), .user("What is the weather in Paris, France?"), ], @@ -147,10 +147,8 @@ struct Gemma4_12BTests { let chat: [Chat.Message] = [ .system(""" - <|think|> You are a helpful assistant that can provide stock prices. When asked for a stock price, you must use the get_stock_price tool. - <|think|> """), .user("What is the price of AAPL?"), ] @@ -166,7 +164,6 @@ struct Gemma4_12BTests { let toolCall1 = try #require(toolCallsOpt1?.first) Swift.print("\(reasoning1 ?? "")\n\(text1 ?? "")") - #expect(reasoning1 != nil) #expect(text1 == nil) #expect(toolCall1.function.name == "get_stock_price") #expect(toolCall1.function.arguments["symbol"] == .string("AAPL")) @@ -194,10 +191,8 @@ struct Gemma4_12BTests { let chat: [Chat.Message] = [ .system(""" - <|think|> You are an email assistant. When asked to read an email, call mail_read exactly once. After the tool result is provided, reply with the email subject exactly and do not call tools again. - <|think|> """), .user("Read email 158348 from account me@example.com in mailbox INBOX."), ] @@ -213,7 +208,6 @@ struct Gemma4_12BTests { let toolCall1 = try #require(toolCallsOpt1?.first) Swift.print("\(reasoning1 ?? "")\n\(text1 ?? "")") - #expect(reasoning1 != nil) #expect(text1 == nil) #expect(toolCall1.function.name == "mail_read") #expect(toolCall1.function.arguments["account"] == .string("me@example.com")) diff --git a/Tests/SHLLMTests/Models/Gemma4-E2BTests.swift b/Tests/SHLLMTests/Models/Gemma4-E2BTests.swift index f77337c..f2275ca 100644 --- a/Tests/SHLLMTests/Models/Gemma4-E2BTests.swift +++ b/Tests/SHLLMTests/Models/Gemma4-E2BTests.swift @@ -104,7 +104,7 @@ struct Gemma4_E2BTests { let input = UserInput( chat: [ .system( - "You are a weather assistant who must use the get_current_weather tool to fetch weather data for any location the user asks about.<|think|>" + "You are a weather assistant who must use the get_current_weather tool to fetch weather data for any location the user asks about." ), .user("What is the weather in Paris, France?"), ], @@ -147,10 +147,8 @@ struct Gemma4_E2BTests { let chat: [Chat.Message] = [ .system(""" - <|think|> You are a helpful assistant that can provide stock prices. When asked for a stock price, you must use the get_stock_price tool. - <|think|> """), .user("What is the price of AAPL?"), ] @@ -166,7 +164,6 @@ struct Gemma4_E2BTests { let toolCall1 = try #require(toolCallsOpt1?.first) Swift.print("\(reasoning1 ?? "")\n\(text1 ?? "")") - #expect(reasoning1 != nil) #expect(text1 == nil) #expect(toolCall1.function.name == "get_stock_price") #expect(toolCall1.function.arguments["symbol"] == .string("AAPL")) @@ -194,10 +191,8 @@ struct Gemma4_E2BTests { let chat: [Chat.Message] = [ .system(""" - <|think|> You are an email assistant. When asked to read an email, call mail_read exactly once. After the tool result is provided, reply with the email subject exactly and do not call tools again. - <|think|> """), .user("Read email 158348 from account me@example.com in mailbox INBOX."), ] @@ -213,7 +208,6 @@ struct Gemma4_E2BTests { let toolCall1 = try #require(toolCallsOpt1?.first) Swift.print("\(reasoning1 ?? "")\n\(text1 ?? "")") - #expect(reasoning1 != nil) #expect(text1 == nil) #expect(toolCall1.function.name == "mail_read") #expect(toolCall1.function.arguments["account"] == .string("me@example.com")) diff --git a/Tests/SHLLMTests/Models/Gemma4-E4BTests.swift b/Tests/SHLLMTests/Models/Gemma4-E4BTests.swift index 003ecaa..fc495f5 100644 --- a/Tests/SHLLMTests/Models/Gemma4-E4BTests.swift +++ b/Tests/SHLLMTests/Models/Gemma4-E4BTests.swift @@ -104,7 +104,7 @@ struct Gemma4_E4BTests { let input = UserInput( chat: [ .system( - "You are a weather assistant who must use the get_current_weather tool to fetch weather data for any location the user asks about.<|think|>" + "You are a weather assistant who must use the get_current_weather tool to fetch weather data for any location the user asks about." ), .user("What is the weather in Paris, France?"), ], @@ -147,10 +147,8 @@ struct Gemma4_E4BTests { let chat: [Chat.Message] = [ .system(""" - <|think|> You are a helpful assistant that can provide stock prices. When asked for a stock price, you must use the get_stock_price tool. - <|think|> """), .user("What is the price of AAPL?"), ] @@ -166,7 +164,6 @@ struct Gemma4_E4BTests { let toolCall1 = try #require(toolCallsOpt1?.first) Swift.print("\(reasoning1 ?? "")\n\(text1 ?? "")") - #expect(reasoning1 != nil) #expect(text1 == nil) #expect(toolCall1.function.name == "get_stock_price") #expect(toolCall1.function.arguments["symbol"] == .string("AAPL")) @@ -194,10 +191,8 @@ struct Gemma4_E4BTests { let chat: [Chat.Message] = [ .system(""" - <|think|> You are an email assistant. When asked to read an email, call mail_read exactly once. After the tool result is provided, reply with the email subject exactly and do not call tools again. - <|think|> """), .user("Read email 158348 from account me@example.com in mailbox INBOX."), ] @@ -213,7 +208,6 @@ struct Gemma4_E4BTests { let toolCall1 = try #require(toolCallsOpt1?.first) Swift.print("\(reasoning1 ?? "")\n\(text1 ?? "")") - #expect(reasoning1 != nil) #expect(text1 == nil) #expect(toolCall1.function.name == "mail_read") #expect(toolCall1.function.arguments["account"] == .string("me@example.com"))