diff --git a/Applications/MLXChatExample/Services/MLXService.swift b/Applications/MLXChatExample/Services/MLXService.swift
index b7a437e8..b9049970 100644
--- a/Applications/MLXChatExample/Services/MLXService.swift
+++ b/Applications/MLXChatExample/Services/MLXService.swift
@@ -131,7 +131,7 @@ class MLXService {
             chat: chat, processing: .init(resize: .init(width: 1024, height: 1024)))
 
         // Generate response using the model
-        return try await modelContainer.perform { (context: ModelContext) in
+        return try await modelContainer.perform { context in
             let lmInput = try await context.processor.prepare(input: userInput)
             // Set temperature for response randomness (0.7 provides good balance)
             let parameters = GenerateParameters(temperature: 0.7)
diff --git a/Tools/llm-tool/LLMTool.swift b/Tools/llm-tool/LLMTool.swift
index 9c95e335..a424a58c 100644
--- a/Tools/llm-tool/LLMTool.swift
+++ b/Tools/llm-tool/LLMTool.swift
@@ -196,7 +196,7 @@ struct GenerateArguments: ParsableArguments, Sendable {
     }
 
     func prepare(
-        _ context: inout ModelContext
+        _ context: inout TrainableModelContext
     ) {
         if let extraEosToken {
             context.configuration.extraEOSTokens.insert(extraEosToken)
@@ -204,7 +204,7 @@ struct GenerateArguments: ParsableArguments, Sendable {
     }
 
     func generate(
-        input: LMInput, context: ModelContext
+        input: LMInput, context: ModelContextProviding
     ) async throws -> (GenerateCompletionInfo, String) {
         var output = ""
         for await item in try MLXLMCommon.generate(