10 files changed
+627
-71
lib/features/on_device/data/on_device_mlx_service.dart
+148
-23
@@ -1,4 +1,5 @@
1
import 'dart:async';
2
+import 'dart:convert';
3
import 'dart:io';
4
5
import 'package:apple_mlx/apple_mlx.dart';
@@ -17,11 +18,28 @@ import '../../chat/data/chat_service.dart';
18
import 'models/on_device_model.dart';
19
20
class OnDeviceMlxService {
20
- OnDeviceMlxService(this._dio);
21
+ OnDeviceMlxService(
22
+ this._dio, {
23
+ this.modelsDirectoryProvider,
24
+ AppleMlxLlmManager Function()? llmManagerFactory,
25
+ bool? nativeInferenceAvailable,
26
+ bool? isIOS,
27
+ }) : _llmManagerFactory = llmManagerFactory ?? AppleMlxLlmManager.new,
28
+ _nativeInferenceAvailable =
29
+ nativeInferenceAvailable ??
30
+ AppleMlxCapabilities.nativeLlmInferenceAvailable,
31
+ _isIOS = isIOS ?? Platform.isIOS;
32
+
33
+ static const String _downloadCompleteFileName = '.download_complete';
34
35
final Dio _dio;
36
+ final Future<Directory> Function()? modelsDirectoryProvider;
37
+ final AppleMlxLlmManager Function() _llmManagerFactory;
38
+ final bool _nativeInferenceAvailable;
39
+ final bool _isIOS;
40
AppleMlxLlmManager? _llmManager;
41
String? _currentModelId;
42
+ bool _currentModelSupportsThinking = false;
43
bool _isDisposed = false;
44
final Map<String, CancelToken> _activeDownloads = {};
45
@@ -30,8 +48,16 @@ class OnDeviceMlxService {
48
49
/// Returns the base directory where MLX models are stored.
50
Future<Directory> get _modelsDirectory async {
33
- final appDir = await getApplicationDocumentsDirectory();
34
- final mlxDir = Directory(p.join(appDir.path, 'models', 'mlx'));
51
+ final override = modelsDirectoryProvider;
52
+ final mlxDir = override != null
53
+ ? await override()
54
+ : Directory(
55
+ p.join(
56
+ (await getApplicationDocumentsDirectory()).path,
57
+ 'models',
58
+ 'mlx',
59
+ ),
60
+ );
61
if (!await mlxDir.exists()) {
62
await mlxDir.create(recursive: true);
63
}
@@ -51,15 +77,33 @@ class OnDeviceMlxService {
77
final dir = await getModelDirectory(modelId);
78
if (!await dir.exists()) return false;
79
80
+ final completionMarker = File(
81
+ p.join(dir.path, _downloadCompleteFileName),
82
+ );
83
+ if (!await completionMarker.exists()) return false;
84
+
85
+ final markerData = jsonDecode(await completionMarker.readAsString());
86
+ if (markerData is! Map || markerData['files'] is! List) return false;
87
+ final manifestFiles = (markerData['files'] as List)
88
+ .whereType<String>()
89
+ .toList(growable: false);
90
+ if (manifestFiles.isEmpty) return false;
91
+
92
+ for (final relativePath in manifestFiles) {
93
+ final filePath = p.normalize(p.join(dir.path, relativePath));
94
+ if (!p.isWithin(dir.path, filePath) || !await File(filePath).exists()) {
95
+ return false;
96
+ }
97
+ }
98
+
99
final config = File(p.join(dir.path, 'config.json'));
100
if (!await config.exists()) return false;
101
57
- final files = await dir.list().toList();
58
- final hasWeights = files.any(
59
- (f) =>
60
- f.path.endsWith('.safetensors') ||
61
- f.path.endsWith('.npz') ||
62
- f.path.endsWith('.bin'),
102
+ final hasWeights = manifestFiles.any(
103
+ (path) =>
104
+ path.endsWith('.safetensors') ||
105
+ path.endsWith('.npz') ||
106
+ path.endsWith('.bin'),
107
);
108
return hasWeights;
109
} catch (e) {
@@ -70,7 +114,7 @@ class OnDeviceMlxService {
114
115
/// Returns a set of installed MLX model IDs.
116
Future<Set<String>> getInstalledModelIds() async {
73
- if (!Platform.isIOS && !Platform.isMacOS) {
117
+ if (!_nativeInferenceAvailable || !_isIOS) {
118
return const <String>{};
119
}
120
try {
@@ -99,9 +143,17 @@ class OnDeviceMlxService {
143
void Function(int receivedBytes, int totalBytes)? onProgress,
144
CancelToken? cancelToken,
145
}) async {
102
- if (!Platform.isIOS && !Platform.isMacOS) {
146
+ _ensureNativeInferenceAvailable();
147
+ if (!_isIOS) {
148
throw UnsupportedError(
104
- 'Apple MLX models are only supported on Apple Silicon (iOS/macOS).',
149
+ 'Apple MLX text generation is only supported on iOS.',
150
+ );
151
+ }
152
+ if (model.runtime != OnDeviceModelRuntime.mlx) {
153
+ throw ArgumentError.value(
154
+ model.runtime,
155
+ 'model.runtime',
156
+ 'Expected an Apple MLX model',
157
);
158
}
159
@@ -109,6 +161,12 @@ class OnDeviceMlxService {
161
if (!await targetDir.exists()) {
162
await targetDir.create(recursive: true);
163
}
164
+ final completionMarker = File(
165
+ p.join(targetDir.path, _downloadCompleteFileName),
166
+ );
167
+ if (await completionMarker.exists()) {
168
+ await completionMarker.delete();
169
+ }
170
171
// Extract repo ID from huggingFaceUrl if needed
172
final repoId = _extractRepoId(model.huggingFaceUrl);
@@ -124,7 +182,8 @@ class OnDeviceMlxService {
182
};
183
184
// 1. Fetch file list from Hugging Face API
127
- final apiUrl = 'https://huggingface.co/api/models/$repoId/tree/main';
185
+ final apiUrl =
186
+ 'https://huggingface.co/api/models/$repoId/tree/main?recursive=true';
187
List<String> filesToDownload = [];
188
int totalModelBytes = model.fileSizeBytes;
189
@@ -177,10 +236,19 @@ class OnDeviceMlxService {
236
final fileSizes = <String, int>{};
237
238
for (final fileName in filesToDownload) {
180
- if (effectiveCancelToken.isCancelled) break;
239
+ if (effectiveCancelToken.isCancelled) {
240
+ throw StateError('MLX model download was cancelled.');
241
+ }
242
243
final fileUrl = 'https://huggingface.co/$repoId/resolve/main/$fileName';
183
- final savePath = p.join(targetDir.path, fileName);
244
+ final savePath = p.normalize(p.join(targetDir.path, fileName));
245
+ if (!p.isWithin(targetDir.path, savePath)) {
246
+ throw StateError('Invalid model file path returned by Hugging Face.');
247
+ }
248
+ final parent = Directory(p.dirname(savePath));
249
+ if (!await parent.exists()) {
250
+ await parent.create(recursive: true);
251
+ }
252
int fileDownloadedBefore = fileSizes[fileName] ?? 0;
253
254
await _dio.download(
@@ -202,6 +270,20 @@ class OnDeviceMlxService {
270
);
271
}
272
273
+ if (effectiveCancelToken.isCancelled) {
274
+ throw StateError('MLX model download was cancelled.');
275
+ }
276
+ for (final fileName in filesToDownload) {
277
+ if (!await File(p.join(targetDir.path, fileName)).exists()) {
278
+ throw StateError('MLX model download is missing $fileName.');
279
+ }
280
+ }
281
+ await completionMarker.writeAsString(
282
+ jsonEncode({'repository': repoId, 'files': filesToDownload}),
283
+ flush: true,
284
+ );
285
+ onProgress?.call(totalModelBytes, totalModelBytes);
286
+
287
Log.info(
288
'MLX model ${model.id} successfully downloaded to ${targetDir.path}',
289
);
@@ -240,6 +322,20 @@ class OnDeviceMlxService {
322
if (_isDisposed) {
323
throw StateError('OnDeviceMlxService has been disposed');
324
}
325
+ _ensureNativeInferenceAvailable();
326
+ if (!_isIOS) {
327
+ throw UnsupportedError(
328
+ 'Apple MLX text generation is only supported on iOS.',
329
+ );
330
+ }
331
+ if (model.runtime != OnDeviceModelRuntime.mlx) {
332
+ throw ArgumentError.value(
333
+ model.runtime,
334
+ 'model.runtime',
335
+ 'Expected an Apple MLX model',
336
+ );
337
+ }
338
+ if (_currentModelId == model.id && isLoaded) return;
339
340
final dir = await getModelDirectory(model.id);
341
if (!await isModelInstalled(model.id)) {
@@ -251,7 +347,7 @@ class OnDeviceMlxService {
347
await unloadModel();
348
349
Log.info('Loading MLX model ${model.id} from ${dir.path}');
254
- final manager = AppleMlxLlmManager();
350
+ final manager = _llmManagerFactory();
351
await manager.loadAndStartServer(
352
modelPath: dir.path,
353
modelId: _extractRepoId(model.huggingFaceUrl),
@@ -260,16 +356,20 @@ class OnDeviceMlxService {
356
357
_llmManager = manager;
358
_currentModelId = model.id;
359
+ _currentModelSupportsThinking = model.supportsThinking;
360
Log.info('MLX model ${model.id} loaded successfully');
361
}
362
363
/// Unloads the active MLX model.
364
Future<void> unloadModel() async {
268
- if (_llmManager != null) {
269
- Log.info('Unloading MLX model $_currentModelId');
270
- await _llmManager!.dispose();
271
- _llmManager = null;
272
- _currentModelId = null;
365
+ final manager = _llmManager;
366
+ final modelId = _currentModelId;
367
+ _llmManager = null;
368
+ _currentModelId = null;
369
+ _currentModelSupportsThinking = false;
370
+ if (manager != null) {
371
+ Log.info('Unloading MLX model $modelId');
372
+ await manager.dispose();
373
}
374
}
375
@@ -288,6 +388,16 @@ class OnDeviceMlxService {
388
yield const ChatResponse(type: ChatResponseType.done);
389
return;
390
}
391
+ if (_currentModelId != modelId) {
392
+ yield ChatResponse(
393
+ type: ChatResponseType.error,
394
+ content:
395
+ 'MLX model $modelId is not loaded. The active model is '
396
+ '${_currentModelId ?? 'unknown'}.',
397
+ );
398
+ yield const ChatResponse(type: ChatResponseType.done);
399
+ return;
400
+ }
401
402
final formattedMessages = <Map<String, String>>[];
403
@@ -299,6 +409,9 @@ class OnDeviceMlxService {
409
}
410
411
for (final m in messages) {
412
+ if (m.role == MessageRole.system && m.content == params.systemPrompt) {
413
+ continue;
414
+ }
415
final role = m.role == MessageRole.user
416
? 'user'
417
: m.role == MessageRole.assistant
@@ -307,12 +420,18 @@ class OnDeviceMlxService {
420
formattedMessages.add({'role': role, 'content': m.content});
421
}
422
423
+ while (formattedMessages.isNotEmpty &&
424
+ formattedMessages.last['role'] == 'assistant' &&
425
+ formattedMessages.last['content']!.isEmpty) {
426
+ formattedMessages.removeLast();
427
+ }
428
+
429
try {
430
final stream = manager.generateStream(
431
messages: formattedMessages,
432
temperature: params.temperature,
433
maxTokens: params.maxTokens,
315
- thinking: true,
434
+ thinking: _currentModelSupportsThinking,
435
);
436
437
await for (final delta in stream) {
@@ -325,7 +444,7 @@ class OnDeviceMlxService {
444
delta.reasoningContent!.isNotEmpty) {
445
yield ChatResponse(
446
type: ChatResponseType.reasoning,
328
- content: delta.reasoningContent!,
447
+ reasoningContent: delta.reasoningContent!,
448
);
449
}
450
@@ -355,6 +474,12 @@ class OnDeviceMlxService {
474
unawaited(unloadModel());
475
}
476
477
+ void _ensureNativeInferenceAvailable() {
478
+ if (!_nativeInferenceAvailable) {
479
+ throw UnsupportedError(AppleMlxCapabilities.nativeLlmUnavailableReason);
480
+ }
481
+ }
482
+
483
String _extractRepoId(String url) {
484
if (!url.startsWith('http')) return url;
485
final uri = Uri.tryParse(url);
lib/features/on_device/providers/on_device_providers.dart
+9
-3
@@ -1,5 +1,6 @@
1
import 'dart:io';
2
3
+import 'package:apple_mlx/apple_mlx.dart';
4
import 'package:dio/dio.dart';
5
import 'package:flutter_gemma/flutter_gemma.dart' show PreferredBackend;
6
import 'package:flutter_riverpod/flutter_riverpod.dart';
@@ -60,11 +61,16 @@ final onDeviceEngineProvider =
61
return OnDeviceEngineNotifier();
62
});
63
64
+List<OnDeviceModel> availableCuratedMlxModels({required bool isIOS}) {
65
+ if (!isIOS || !AppleMlxCapabilities.nativeLlmInferenceAvailable) {
66
+ return const <OnDeviceModel>[];
67
+ }
68
+ return OnDeviceModel.curatedMlxModels;
69
+}
70
+
71
final onDeviceModelsProvider = Provider<List<OnDeviceModel>>((ref) {
72
final imported = ref.watch(importedGgufModelsProvider);
65
- final mlxModels = Platform.isIOS
66
- ? OnDeviceModel.curatedMlxModels
67
- : const <OnDeviceModel>[];
73
+ final mlxModels = availableCuratedMlxModels(isIOS: Platform.isIOS);
74
return [...mlxModels, ...OnDeviceModel.curatedModels, ...imported];
75
});
76
packages/apple_mlx/README.md
+16
-5
@@ -1,16 +1,24 @@
1
# Apple MLX Support (`apple_mlx`)
2
3
-Local package integrating Apple's **MLX** ecosystem (`lib_mlx`, `speech_mlx`, `mlx_audio`, and `mlx-c`) with Dart & Flutter to run LLM Text Generation, Speech-to-Text, and Text-to-Speech natively on Apple Silicon (iOS and macOS).
3
+Local package collecting experimental Dart wrappers around `lib_mlx`,
4
+`speech_mlx`, and `mlx_audio`.
5
+
6
+> **Development status:** Native LLM inference is disabled in LocalMind. The
7
+> current `lib_mlx` 0.1.1 dependency ships a mock model core that returns stub
8
+> responses; it does not execute downloaded MLX weights. Its bridge targets
9
+> iOS only. The speech dependencies target macOS Apple Silicon and are not
10
+> integrated into the LocalMind mobile app.
11
12
## Capabilities
13
14
### 1. LLM Text Generation (`AppleMlxLlmManager` & `lib_mlx`)
8
-- **Native Apple Silicon Inference**: Runs quantized MLX models (Gemma, Llama 3.2, Qwen 2.5, SmolLM2) directly on device via Unified Memory and Metal.
9
-- **OpenAI-Compatible Local Endpoint**: Zero-overhead localhost HTTP/SSE streaming client.
10
-- **Reasoning / Thinking Support**: Separates `<thinking>` tokens from final assistant content.
11
-- **Tool Calling Support**: Function and tool calling protocol.
15
+- **Prototype only**: Exercises the native lifecycle and local OpenAI-compatible
16
+ HTTP/SSE contract against the upstream mock core.
17
+- **Not user-facing**: `AppleMlxCapabilities.nativeLlmInferenceAvailable` stays
18
+ `false` until a real, device-validated backend replaces the mock.
19
20
### 2. Speech-To-Text (`AppleMlxSpeechManager` & `speech_mlx` / `mlx_audio`)
21
+- **macOS only**: These dependencies do not provide iOS speech inference.
22
- **Parakeet-TDT & Whisper**: High-accuracy local speech recognition.
23
- **Isolate Offloading**: Dedicated worker isolates keep audio processing and resampling off the UI isolate.
24
@@ -27,6 +35,9 @@ Local package integrating Apple's **MLX** ecosystem (`lib_mlx`, `speech_mlx`, `m
35
36
### LLM Text Generation
37
38
+The following example is for bridge development and contract testing only. It
39
+does not perform real inference with `lib_mlx` 0.1.1.
40
+
41
```dart
42
import 'package:apple_mlx/apple_mlx.dart';
43
packages/apple_mlx/lib/apple_mlx.dart
+1
@@ -8,6 +8,7 @@ export 'package:speech_core/speech_core.dart';
8
export 'package:audio_core/audio_core.dart';
9
10
// LocalMind MLX Suite
11
+export 'src/mlx_capabilities.dart';
12
export 'src/mlx_models.dart';
13
export 'src/mlx_speech_manager.dart';
14
export 'src/mlx_llm_manager.dart';
packages/apple_mlx/lib/src/mlx_capabilities.dart
new
+15
@@ -0,0 +1,15 @@
1
+/// Runtime capabilities of the native packages wrapped by `apple_mlx`.
2
+///
3
+/// Keep these flags conservative. A capability should only become available
4
+/// after the native dependency performs real inference on a supported device.
5
+abstract final class AppleMlxCapabilities {
6
+ /// `lib_mlx` 0.1.1 ships a mock native model core, not MLX inference.
7
+ ///
8
+ /// LocalMind uses this flag to keep the MLX download and chat path out of the
9
+ /// production catalog until the native backend is replaced by a real one.
10
+ static const bool nativeLlmInferenceAvailable = false;
11
+
12
+ static const String nativeLlmUnavailableReason =
13
+ 'Apple MLX text generation is unavailable because the current native '
14
+ 'lib_mlx backend is a mock scaffold and does not run model inference.';
15
+}
packages/apple_mlx/lib/src/mlx_llm_manager.dart
+77
-40
@@ -1,6 +1,10 @@
1
import 'dart:async';
2
import 'package:lib_mlx/lib_mlx.dart';
3
4
+import 'mlx_capabilities.dart';
5
+
6
+typedef AppleMlxOpenAiClientFactory = LibMlxOpenAiClient Function(Uri baseUri);
7
+
8
/// Represents streaming chunk delta from an MLX LLM generation.
9
class AppleMlxChatDelta {
10
/// Regular content delta chunk.
@@ -23,15 +27,23 @@ class AppleMlxChatDelta {
27
});
28
}
29
26
-/// High-level manager for running MLX LLM text generation models natively on Apple devices.
30
+/// High-level client for the native `lib_mlx` lifecycle and localhost server.
31
+///
32
+/// Check [AppleMlxCapabilities.nativeLlmInferenceAvailable] before exposing
33
+/// this manager as an inference backend.
34
class AppleMlxLlmManager {
35
final LibMlxRuntime _runtime;
36
+ final AppleMlxOpenAiClientFactory _clientFactory;
37
LibMlxOpenAiClient? _client;
38
MlxModelHandle? _handle;
39
MlxServerInfo? _serverInfo;
40
33
- AppleMlxLlmManager({LibMlxRuntime? runtime})
34
- : _runtime = runtime ?? const LibMlxRuntime();
41
+ AppleMlxLlmManager({
42
+ LibMlxRuntime? runtime,
43
+ AppleMlxOpenAiClientFactory? clientFactory,
44
+ }) : _runtime = runtime ?? const LibMlxRuntime(),
45
+ _clientFactory =
46
+ clientFactory ?? ((baseUri) => LibMlxOpenAiClient(baseUri: baseUri));
47
48
bool get isModelLoaded => _handle != null;
49
bool get isServerRunning => _serverInfo != null;
@@ -47,7 +59,6 @@ class AppleMlxLlmManager {
59
int port = 0,
60
int queueLimit = 1,
61
}) async {
50
- // Unload existing session if present
62
await dispose();
63
64
final config = MlxModelConfig(
@@ -57,18 +68,29 @@ class AppleMlxLlmManager {
68
lazyEncoders: lazyEncoders,
69
);
70
60
- _handle = await _runtime.loadModel(config);
61
-
62
- final serverConfig = MlxServerConfig(
63
- port: port,
64
- modelId: modelId,
65
- queueLimit: queueLimit,
66
- );
71
+ try {
72
+ _handle = await _runtime.loadModel(config);
73
68
- _serverInfo = await _runtime.startServer(_handle!, config: serverConfig);
69
- _client = LibMlxOpenAiClient(baseUri: _serverInfo!.uri);
74
+ final serverConfig = MlxServerConfig(
75
+ port: port,
76
+ modelId: modelId,
77
+ queueLimit: queueLimit,
78
+ );
79
71
- return _serverInfo!;
80
+ final serverInfo = await _runtime.startServer(
81
+ _handle!,
82
+ config: serverConfig,
83
+ );
84
+ _serverInfo = serverInfo;
85
+ _client = _clientFactory(serverInfo.uri);
86
+ return serverInfo;
87
+ } catch (_) {
88
+ // Loading succeeds before the server is started. If either server start
89
+ // or client construction fails, release that resident native handle so
90
+ // the next attempt does not leak model memory.
91
+ await dispose();
92
+ rethrow;
93
+ }
94
}
95
96
/// Streams chat completions using Server-Sent Events (SSE).
@@ -98,27 +120,35 @@ class AppleMlxLlmManager {
120
121
await for (final event in sseStream) {
122
if (event.done) {
101
- yield const AppleMlxChatDelta(isDone: true);
123
break;
124
}
125
105
- final choices = event.data?['choices'] as List?;
106
- final choice = choices?.firstOrNull as Map?;
107
- final delta = choice?['delta'] as Map?;
108
-
109
- if (delta != null) {
110
- final content = delta['content'] as String?;
111
- final reasoning = delta['reasoning_content'] as String?;
112
-
113
- if (content != null || reasoning != null) {
114
- yield AppleMlxChatDelta(
115
- content: content,
116
- reasoningContent: reasoning,
117
- rawData: event.data,
118
- );
119
- }
126
+ final choices = event.data?['choices'];
127
+ if (choices is! List || choices.isEmpty) continue;
128
+
129
+ final choice = choices.first;
130
+ if (choice is! Map) continue;
131
+
132
+ final delta = choice['delta'];
133
+ if (delta is! Map) continue;
134
+
135
+ final rawContent = delta['content'];
136
+ final rawReasoning = delta['reasoning_content'];
137
+ final content = rawContent is String ? rawContent : null;
138
+ final reasoning = rawReasoning is String ? rawReasoning : null;
139
+
140
+ if (content != null || reasoning != null) {
141
+ yield AppleMlxChatDelta(
142
+ content: content,
143
+ reasoningContent: reasoning,
144
+ rawData: event.data,
145
+ );
146
}
147
}
148
+
149
+ // Some HTTP servers close a valid SSE response without a final [DONE]
150
+ // sentinel. Always terminate the higher-level stream consistently.
151
+ yield const AppleMlxChatDelta(isDone: true);
152
}
153
154
/// Generates non-streaming completion response.
@@ -143,12 +173,17 @@ class AppleMlxLlmManager {
173
};
174
175
final response = await client.chatCompletions(request);
146
- final choices = response['choices'] as List?;
147
- if (choices != null && choices.isNotEmpty) {
148
- final first = choices.first as Map?;
149
- final message = first?['message'] as Map?;
150
- return message?['content'] as String? ?? '';
151
- }
176
+ final choices = response['choices'];
177
+ if (choices is! List || choices.isEmpty) return '';
178
+
179
+ final first = choices.first;
180
+ if (first is! Map) return '';
181
+
182
+ final message = first['message'];
183
+ if (message is! Map) return '';
184
+
185
+ final content = message['content'];
186
+ if (content is String) return content;
187
return '';
188
}
189
@@ -161,10 +196,14 @@ class AppleMlxLlmManager {
196
197
/// Stops server, unloads model, and releases memory.
198
Future<void> dispose() async {
164
- _client?.close();
199
+ final client = _client;
200
+ final handle = _handle;
201
+
202
_client = null;
203
+ _handle = null;
204
+ _serverInfo = null;
205
+ client?.close(force: true);
206
167
- final handle = _handle;
207
if (handle != null) {
208
try {
209
await _runtime.stopServer(handle);
@@ -172,8 +211,6 @@ class AppleMlxLlmManager {
211
try {
212
await _runtime.unloadModel(handle);
213
} catch (_) {}
175
- _handle = null;
214
}
177
- _serverInfo = null;
215
}
216
}
packages/apple_mlx/lib/src/mlx_models.dart
+2
@@ -3,6 +3,8 @@ class AppleMlxModelIds {
3
const AppleMlxModelIds._();
4
5
// LLM / Text Generation Models (Hugging Face mlx-community)
6
+ static const String gemma4E2b = 'mlx-community/gemma-4-e2b-it-4bit';
7
+ static const String gemma4E4b = 'mlx-community/gemma-4-e4b-it-4bit';
8
static const String gemma22b = 'mlx-community/gemma-2-2b-it-4bit';
9
static const String llama321b = 'mlx-community/Llama-3.2-1B-Instruct-4bit';
10
static const String llama323b = 'mlx-community/Llama-3.2-3B-Instruct-4bit';
packages/apple_mlx/test/apple_mlx_test.dart
+122
@@ -36,9 +36,131 @@ void main() {
36
});
37
38
test('exposes popular MLX model identifiers', () {
39
+ expect(AppleMlxModelIds.gemma4E2b, isNotEmpty);
40
expect(AppleMlxModelIds.gemma22b, isNotEmpty);
41
expect(AppleMlxModelIds.llama321b, isNotEmpty);
42
expect(AppleMlxModelIds.qwen2505b, isNotEmpty);
43
});
44
+
45
+ test('cleans up a loaded model when server startup fails', () async {
46
+ final runtime = _FakeRuntime(failServerStart: true);
47
+ final llm = AppleMlxLlmManager(runtime: runtime);
48
+
49
+ await expectLater(
50
+ llm.loadAndStartServer(modelPath: '/models/test'),
51
+ throwsStateError,
52
+ );
53
+
54
+ expect(runtime.loadCount, 1);
55
+ expect(runtime.unloadCount, 1);
56
+ expect(llm.isModelLoaded, isFalse);
57
+ expect(llm.isServerRunning, isFalse);
58
+ });
59
+
60
+ test('emits one terminal delta when SSE includes DONE', () async {
61
+ final client = _FakeClient([
62
+ const MlxSseEvent(
63
+ data: {
64
+ 'choices': [
65
+ {
66
+ 'delta': {'content': 'answer', 'reasoning_content': 'thought'},
67
+ },
68
+ ],
69
+ },
70
+ ),
71
+ const MlxSseEvent(done: true),
72
+ ]);
73
+ final llm = AppleMlxLlmManager(
74
+ runtime: _FakeRuntime(),
75
+ clientFactory: (_) => client,
76
+ );
77
+ await llm.loadAndStartServer(modelPath: '/models/test');
78
+
79
+ final deltas = await llm
80
+ .generateStream(
81
+ messages: const [
82
+ {'role': 'user', 'content': 'hello'},
83
+ ],
84
+ )
85
+ .toList();
86
+
87
+ expect(deltas, hasLength(2));
88
+ expect(deltas.first.content, 'answer');
89
+ expect(deltas.first.reasoningContent, 'thought');
90
+ expect(deltas.last.isDone, isTrue);
91
+ expect(deltas.where((delta) => delta.isDone), hasLength(1));
92
+ });
93
+
94
+ test('emits a terminal delta when SSE closes without DONE', () async {
95
+ final client = _FakeClient(const []);
96
+ final llm = AppleMlxLlmManager(
97
+ runtime: _FakeRuntime(),
98
+ clientFactory: (_) => client,
99
+ );
100
+ await llm.loadAndStartServer(modelPath: '/models/test');
101
+
102
+ final deltas = await llm
103
+ .generateStream(
104
+ messages: const [
105
+ {'role': 'user', 'content': 'hello'},
106
+ ],
107
+ )
108
+ .toList();
109
+
110
+ expect(deltas, hasLength(1));
111
+ expect(deltas.single.isDone, isTrue);
112
+ });
113
+ });
114
+
115
+ test('native LLM inference remains disabled for the mock backend', () {
116
+ expect(AppleMlxCapabilities.nativeLlmInferenceAvailable, isFalse);
117
+ expect(AppleMlxCapabilities.nativeLlmUnavailableReason, contains('mock'));
118
});
119
}
120
+
121
+class _FakeRuntime extends LibMlxRuntime {
122
+ _FakeRuntime({this.failServerStart = false});
123
+
124
+ final bool failServerStart;
125
+ int loadCount = 0;
126
+ int unloadCount = 0;
127
+
128
+ @override
129
+ Future<MlxModelHandle> loadModel(MlxModelConfig config) async {
130
+ loadCount++;
131
+ return MlxModelHandle(value: loadCount, modelId: config.modelId);
132
+ }
133
+
134
+ @override
135
+ Future<MlxServerInfo> startServer(
136
+ MlxModelHandle handle, {
137
+ MlxServerConfig config = const MlxServerConfig(),
138
+ }) async {
139
+ if (failServerStart) throw StateError('server failed');
140
+ return MlxServerInfo(
141
+ host: '127.0.0.1',
142
+ port: 12345,
143
+ baseUrl: 'http://127.0.0.1:12345',
144
+ modelId: config.modelId,
145
+ status: 'running',
146
+ );
147
+ }
148
+
149
+ @override
150
+ Future<void> stopServer(MlxModelHandle handle) async {}
151
+
152
+ @override
153
+ Future<void> unloadModel(MlxModelHandle handle) async {
154
+ unloadCount++;
155
+ }
156
+}
157
+
158
+class _FakeClient extends LibMlxOpenAiClient {
159
+ _FakeClient(this.events) : super(baseUri: Uri.parse('http://localhost'));
160
+
161
+ final List<MlxSseEvent> events;
162
+
163
+ @override
164
+ Stream<MlxSseEvent> chatCompletionsStream(Map<String, Object?> request) =>
165
+ Stream.fromIterable(events);
166
+}
test/features/chat/chat_service_routing_test.dart
+19
@@ -8,6 +8,7 @@ import 'package:localmind/features/on_device/data/models/on_device_model.dart';
8
import 'package:localmind/features/on_device/data/on_device_chat_service.dart';
9
import 'package:localmind/features/on_device/data/on_device_gemma_service.dart';
10
import 'package:localmind/features/on_device/data/on_device_llama_service.dart';
11
+import 'package:localmind/features/on_device/data/on_device_mlx_service.dart';
12
import 'package:localmind/features/on_device/providers/on_device_providers.dart';
13
import 'package:localmind/features/servers/data/models/server.dart';
14
import 'package:localmind/features/servers/providers/server_providers.dart';
@@ -38,6 +39,24 @@ void main() {
39
40
expect(service, isA<OnDeviceChatService>());
41
});
42
+
43
+ test('routes an MLX runtime to the MLX chat adapter', () {
44
+ final service = createChatServiceForServer(
45
+ server: _onDeviceServer(),
46
+ dio: Dio(),
47
+ onDeviceGemmaService: OnDeviceGemmaService(),
48
+ onDeviceLlamaService: OnDeviceLlamaService(),
49
+ onDeviceMlxService: OnDeviceMlxService(Dio()),
50
+ loadedOnDeviceRuntime: OnDeviceModelRuntime.mlx,
51
+ );
52
+
53
+ expect(service, isA<OnDeviceMlxChatService>());
54
+ });
55
+ });
56
+
57
+ test('does not expose MLX downloads while the native backend is a mock', () {
58
+ expect(availableCuratedMlxModels(isIOS: true), isEmpty);
59
+ expect(availableCuratedMlxModels(isIOS: false), isEmpty);
60
});
61
62
group('chatServiceProvider', () {
test/features/on_device/on_device_mlx_service_test.dart
new
+218
@@ -0,0 +1,218 @@
1
+import 'dart:convert';
2
+import 'dart:io';
3
+
4
+import 'package:apple_mlx/apple_mlx.dart';
5
+import 'package:dio/dio.dart';
6
+import 'package:flutter_test/flutter_test.dart';
7
+import 'package:localmind/core/models/enums.dart';
8
+import 'package:localmind/features/chat/data/chat_service.dart';
9
+import 'package:localmind/features/chat/data/models/chat_parameters.dart';
10
+import 'package:localmind/features/chat/data/models/message.dart';
11
+import 'package:localmind/features/on_device/data/models/on_device_model.dart';
12
+import 'package:localmind/features/on_device/data/on_device_mlx_service.dart';
13
+import 'package:path/path.dart' as p;
14
+
15
+void main() {
16
+ late Directory modelsDirectory;
17
+
18
+ setUp(() async {
19
+ modelsDirectory = await Directory.systemTemp.createTemp(
20
+ 'localmind_mlx_test_',
21
+ );
22
+ });
23
+
24
+ tearDown(() async {
25
+ if (await modelsDirectory.exists()) {
26
+ await modelsDirectory.delete(recursive: true);
27
+ }
28
+ });
29
+
30
+ test('requires a complete download manifest for installation', () async {
31
+ final service = _service(modelsDirectory: modelsDirectory);
32
+ final model = OnDeviceModel.curatedMlxModels.first;
33
+ final modelDirectory = await service.getModelDirectory(model.id);
34
+ await modelDirectory.create(recursive: true);
35
+ await File(p.join(modelDirectory.path, 'config.json')).writeAsString('{}');
36
+ await File(
37
+ p.join(modelDirectory.path, 'model-00001-of-00002.safetensors'),
38
+ ).writeAsString('weights');
39
+
40
+ expect(await service.isModelInstalled(model.id), isFalse);
41
+
42
+ await File(p.join(modelDirectory.path, '.download_complete')).writeAsString(
43
+ jsonEncode({
44
+ 'files': [
45
+ 'config.json',
46
+ 'model-00001-of-00002.safetensors',
47
+ 'model-00002-of-00002.safetensors',
48
+ ],
49
+ }),
50
+ );
51
+
52
+ expect(await service.isModelInstalled(model.id), isFalse);
53
+
54
+ await File(
55
+ p.join(modelDirectory.path, 'model-00002-of-00002.safetensors'),
56
+ ).writeAsString('weights');
57
+
58
+ expect(await service.isModelInstalled(model.id), isTrue);
59
+ });
60
+
61
+ test('maps MLX reasoning and normalizes the outgoing prompt', () async {
62
+ final manager = _FakeMlxManager();
63
+ final service = _service(
64
+ modelsDirectory: modelsDirectory,
65
+ manager: manager,
66
+ nativeInferenceAvailable: true,
67
+ );
68
+ final model = OnDeviceModel.curatedMlxModels.first;
69
+ await _writeInstalledModel(service, model);
70
+ await service.loadModel(model);
71
+
72
+ final responses = await service
73
+ .sendMessage(
74
+ modelId: model.id,
75
+ messages: [
76
+ _message(MessageRole.system, 'Be concise.'),
77
+ _message(MessageRole.user, 'Hello'),
78
+ _message(MessageRole.assistant, ''),
79
+ ],
80
+ params: ChatParameters.defaults().copyWith(
81
+ systemPrompt: 'Be concise.',
82
+ ),
83
+ )
84
+ .toList();
85
+
86
+ expect(responses.map((response) => response.type), [
87
+ ChatResponseType.reasoning,
88
+ ChatResponseType.message,
89
+ ChatResponseType.done,
90
+ ]);
91
+ expect(responses.first.reasoningContent, 'thinking');
92
+ expect(responses.first.content, isNull);
93
+ expect(responses[1].content, 'answer');
94
+ expect(manager.lastThinking, isFalse);
95
+ expect(manager.lastMessages, [
96
+ {'role': 'system', 'content': 'Be concise.'},
97
+ {'role': 'user', 'content': 'Hello'},
98
+ ]);
99
+ });
100
+
101
+ test('rejects a request for a model other than the loaded model', () async {
102
+ final manager = _FakeMlxManager();
103
+ final service = _service(
104
+ modelsDirectory: modelsDirectory,
105
+ manager: manager,
106
+ nativeInferenceAvailable: true,
107
+ );
108
+ final model = OnDeviceModel.curatedMlxModels.first;
109
+ await _writeInstalledModel(service, model);
110
+ await service.loadModel(model);
111
+
112
+ final responses = await service
113
+ .sendMessage(
114
+ modelId: 'different-model',
115
+ messages: [_message(MessageRole.user, 'Hello')],
116
+ params: ChatParameters.defaults(),
117
+ )
118
+ .toList();
119
+
120
+ expect(responses, hasLength(2));
121
+ expect(responses.first.type, ChatResponseType.error);
122
+ expect(responses.first.content, contains('different-model'));
123
+ expect(responses.last.type, ChatResponseType.done);
124
+ expect(manager.generateCount, 0);
125
+ });
126
+}
127
+
128
+OnDeviceMlxService _service({
129
+ required Directory modelsDirectory,
130
+ _FakeMlxManager? manager,
131
+ bool nativeInferenceAvailable = false,
132
+}) {
133
+ return OnDeviceMlxService(
134
+ Dio(),
135
+ modelsDirectoryProvider: () async => modelsDirectory,
136
+ llmManagerFactory: () => manager ?? _FakeMlxManager(),
137
+ nativeInferenceAvailable: nativeInferenceAvailable,
138
+ isIOS: true,
139
+ );
140
+}
141
+
142
+Future<void> _writeInstalledModel(
143
+ OnDeviceMlxService service,
144
+ OnDeviceModel model,
145
+) async {
146
+ final directory = await service.getModelDirectory(model.id);
147
+ await directory.create(recursive: true);
148
+ await File(p.join(directory.path, 'config.json')).writeAsString('{}');
149
+ await File(
150
+ p.join(directory.path, 'model.safetensors'),
151
+ ).writeAsString('weights');
152
+ await File(p.join(directory.path, '.download_complete')).writeAsString(
153
+ jsonEncode({
154
+ 'files': ['config.json', 'model.safetensors'],
155
+ }),
156
+ );
157
+}
158
+
159
+Message _message(MessageRole role, String content) {
160
+ return Message(
161
+ id: '$role-$content',
162
+ conversationId: 'conversation',
163
+ role: role,
164
+ content: content,
165
+ createdAt: DateTime.utc(2026, 8, 31),
166
+ );
167
+}
168
+
169
+class _FakeMlxManager extends AppleMlxLlmManager {
170
+ bool _isLoaded = false;
171
+ int generateCount = 0;
172
+ bool? lastThinking;
173
+ List<Map<String, String>>? lastMessages;
174
+
175
+ @override
176
+ bool get isModelLoaded => _isLoaded;
177
+
178
+ @override
179
+ Future<MlxServerInfo> loadAndStartServer({
180
+ required String modelPath,
181
+ String modelId = MlxAvailableModel.defaultHuggingFaceModelId,
182
+ bool thinkingEnabled = true,
183
+ bool lazyEncoders = true,
184
+ int port = 0,
185
+ int queueLimit = 1,
186
+ }) async {
187
+ _isLoaded = true;
188
+ return MlxServerInfo(
189
+ host: '127.0.0.1',
190
+ port: 12345,
191
+ baseUrl: 'http://127.0.0.1:12345',
192
+ modelId: modelId,
193
+ status: 'running',
194
+ );
195
+ }
196
+
197
+ @override
198
+ Stream<AppleMlxChatDelta> generateStream({
199
+ required List<Map<String, String>> messages,
200
+ double temperature = 0.7,
201
+ int? maxTokens,
202
+ bool thinking = true,
203
+ }) {
204
+ generateCount++;
205
+ lastMessages = messages;
206
+ lastThinking = thinking;
207
+ return Stream.fromIterable(const [
208
+ AppleMlxChatDelta(reasoningContent: 'thinking'),
209
+ AppleMlxChatDelta(content: 'answer'),
210
+ AppleMlxChatDelta(isDone: true),
211
+ ]);
212
+ }
213
+
214
+ @override
215
+ Future<void> dispose() async {
216
+ _isLoaded = false;
217
+ }
218
+}