diff --git a/Sources/LLM/EspressoLLMEngine.swift b/Sources/LLM/EspressoLLMEngine.swift index a63008e..c413859 100644 --- a/Sources/LLM/EspressoLLMEngine.swift +++ b/Sources/LLM/EspressoLLMEngine.swift @@ -16,6 +16,8 @@ private func aneLMTokenCallback(_ token: Int32, _ context: UnsafeMutableRawPoint } actor EspressoLLMEngine { + static let maximumContextTokens = 2_048 + private final class LoadedModel { let path: String let runtime: OpaquePointer @@ -90,6 +92,18 @@ actor EspressoLLMEngine { .map { Int32(clamping: $0) } guard !promptTokens.isEmpty else { throw EspressoLLMError.runtimeFailure } + let nativeMaxTokens = max(1, maxTokens) + guard Self.requestFitsContextWindow( + promptTokenCount: promptTokens.count, + maxTokens: nativeMaxTokens + ) else { + let requiredCacheSlots = promptTokens.count + max(0, nativeMaxTokens - 1) + throw recordFailure(ANELMNativeError( + "ANE-LM request needs \(requiredCacheSlots) KV-cache slots; " + + "the packaged runtime supports \(Self.maximumContextTokens)" + )) + } + let context = ANELMGenerationContext() context.tokens.reserveCapacity(max(0, maxTokens)) let contextPointer = Unmanaged.passUnretained(context).toOpaque() @@ -103,7 +117,7 @@ actor EspressoLLMEngine { model.runtime, tokens.baseAddress, tokens.count, - Int32(clamping: max(1, maxTokens)), + Int32(clamping: nativeMaxTokens), Float(temperature), 1.2, Int32(clamping: model.samplerVocabularySize), @@ -153,6 +167,16 @@ actor EspressoLLMEngine { _ = try await makeValidatedTokenizer(at: url) } + static func requestFitsContextWindow( + promptTokenCount: Int, + maxTokens: Int + ) -> Bool { + guard promptTokenCount > 0, maxTokens > 0 else { return false } + let generatedCacheSlots = max(0, maxTokens - 1) + guard generatedCacheSlots <= maximumContextTokens else { return false } + return promptTokenCount <= maximumContextTokens - generatedCacheSlots + } + private struct ValidatedTokenizer { let tokenizer: any Tokenizers.Tokenizer let samplerVocabularySize: Int @@ -208,10 +232,14 @@ actor EspressoLLMEngine { } static func formatPrompt(user: String, system: String, modelName: String) -> String { - if modelName.lowercased().contains("qwen") { + let normalizedModelName = modelName.lowercased() + if normalizedModelName.contains("qwen") { + let assistantPrefix = normalizedModelName.contains("qwen3") + ? "<|im_start|>assistant\n\n\n\n\n" + : "<|im_start|>assistant\n" return "<|im_start|>system\n\(system)<|im_end|>\n" + "<|im_start|>user\n\(user)<|im_end|>\n" - + "<|im_start|>assistant\n" + + assistantPrefix } return "System:\n\(system)\n\nUser:\n\(user)\n\nAssistant:\n" } diff --git a/Sources/Processing/TextProcessor+Generation.swift b/Sources/Processing/TextProcessor+Generation.swift index 0a3d4ae..f952712 100644 --- a/Sources/Processing/TextProcessor+Generation.swift +++ b/Sources/Processing/TextProcessor+Generation.swift @@ -48,6 +48,9 @@ extension TextProcessor { temperature: temperature ) }, + prepareForMLXFallback: { + await self.espressoLLM.unload() + }, mlx: { try await self.llm.loadModel(id: options.llmModel) return try await self.llm.generate( @@ -59,9 +62,8 @@ extension TextProcessor { } ) if result.usedMLX { - _ = await espressoLLM.consumeLastFailureMessage() await Self.recordEspressoOutcome(.fallback) - Log.info("[TextProcessor] ANE-LM failed; used the selected MLX model") + Log.info("[TextProcessor] ANE-LM failed; unloaded it and used the selected MLX model") } else { await Self.clearEspressoOutcome() } @@ -69,7 +71,6 @@ extension TextProcessor { } catch is CancellationError { throw CancellationError() } catch let error as EspressoMLXFallbackError { - _ = await espressoLLM.consumeLastFailureMessage() await Self.recordEspressoOutcome(.unavailable) Log.sensitive("[TextProcessor] ANE-LM and MLX fallback failed: \(error.details)") Log.error("[TextProcessor] MLX fallback unavailable") @@ -89,6 +90,7 @@ extension TextProcessor { static func runEspressoWithMLXFallback( fallbackEnabled: Bool = true, espresso: () async throws -> Value, + prepareForMLXFallback: () async -> Void = {}, mlx: () async throws -> Value ) async throws -> (value: Value, usedMLX: Bool) { do { @@ -99,6 +101,8 @@ extension TextProcessor { try Task.checkCancellation() guard fallbackEnabled else { throw error } let espressoFailure = error.localizedDescription + await prepareForMLXFallback() + try Task.checkCancellation() do { let value = try await mlx() try Task.checkCancellation() diff --git a/Sources/Processing/TextProcessor+Models.swift b/Sources/Processing/TextProcessor+Models.swift index 29820e2..44a5620 100644 --- a/Sources/Processing/TextProcessor+Models.swift +++ b/Sources/Processing/TextProcessor+Models.swift @@ -163,10 +163,10 @@ extension TextProcessor { let result = try await Self.runEspressoWithMLXFallback( fallbackEnabled: fallbackToMLXOnEspressoFailure, espresso: { try await self.espressoLLM.loadModel(path: espressoModelPath) }, + prepareForMLXFallback: { await self.espressoLLM.unload() }, mlx: { try await self.llm.loadModel(id: model) } ) if result.usedMLX { - _ = await espressoLLM.consumeLastFailureMessage() return (true, nil, .fallback) } } @@ -176,7 +176,6 @@ extension TextProcessor { } catch let error as EspressoMLXFallbackError { Log.sensitive("[TextProcessor] ANE-LM and MLX warmup failed: \(error.details)") Log.error("[TextProcessor] MLX fallback unavailable during warmup") - _ = await espressoLLM.consumeLastFailureMessage() return (false, EspressoGenerationOutcome.unavailable.message, .unavailable) } catch { Log.error("[TextProcessor] LLM warmup failed: \(error.localizedDescription)") diff --git a/Tests/OpenTypeTests/ANELMRuntimeTests.swift b/Tests/OpenTypeTests/ANELMRuntimeTests.swift index 40dce3e..8afd849 100644 --- a/Tests/OpenTypeTests/ANELMRuntimeTests.swift +++ b/Tests/OpenTypeTests/ANELMRuntimeTests.swift @@ -120,6 +120,41 @@ final class ANELMRuntimeTests: XCTestCase { await XCTAssertThrowsErrorAsync(try await EspressoLLMEngine.validateModelDirectory(at: invalidHeads)) } + func testQwenPromptDisablesThinkingBeforeGeneration() { + let prompt = EspressoLLMEngine.formatPrompt( + user: "Return JSON.", + system: "Do not explain.", + modelName: "Qwen3" + ) + + XCTAssertTrue(prompt.hasSuffix( + "<|im_start|>assistant\n\n\n\n\n" + )) + } + + func testContextWindowGuardReservesGeneratedCacheSlots() { + XCTAssertTrue(EspressoLLMEngine.requestFitsContextWindow( + promptTokenCount: 2_048, + maxTokens: 1 + )) + XCTAssertTrue(EspressoLLMEngine.requestFitsContextWindow( + promptTokenCount: 2_047, + maxTokens: 2 + )) + XCTAssertFalse(EspressoLLMEngine.requestFitsContextWindow( + promptTokenCount: 2_049, + maxTokens: 1 + )) + XCTAssertFalse(EspressoLLMEngine.requestFitsContextWindow( + promptTokenCount: 2_048, + maxTokens: 2 + )) + XCTAssertFalse(EspressoLLMEngine.requestFitsContextWindow( + promptTokenCount: 1, + maxTokens: 2_049 + )) + } + func testRealGenerationLifecycleWhenModelIsProvided() async throws { guard let modelPath = ProcessInfo.processInfo.environment["UTTER_ANE_TEST_MODEL"], !modelPath.isEmpty else { @@ -142,14 +177,24 @@ final class ANELMRuntimeTests: XCTestCase { let requestCount = iterations / lifecycleCount + (lifecycle < iterations % lifecycleCount ? 1 : 0) var lifecycleSamples: [Int] = [] - for _ in 0..")) + if verifiesStructuredCommandOutput { + XCTAssertNotNil(SpokenEditCommandLLMResolver.resolution(from: output)) + } let sample = try residentSizeKB() residentSamples.append(sample) lifecycleSamples.append(sample) diff --git a/Tests/OpenTypeTests/EspressoFallbackTests.swift b/Tests/OpenTypeTests/EspressoFallbackTests.swift index f13dca4..fd3bbe1 100644 --- a/Tests/OpenTypeTests/EspressoFallbackTests.swift +++ b/Tests/OpenTypeTests/EspressoFallbackTests.swift @@ -56,13 +56,38 @@ final class EspressoFallbackTests: XCTestCase { XCTAssertTrue(result.usedMLX) } + func testEspressoFailureReleasesANEStateBeforeMLXFallback() async throws { + var espressoIsLoaded = true + var espressoWasLoadedWhenMLXStarted = true + + let result: (value: String, usedMLX: Bool) = try await TextProcessor.runEspressoWithMLXFallback( + espresso: { throw StubError.espresso }, + prepareForMLXFallback: { + espressoIsLoaded = false + }, + mlx: { + espressoWasLoadedWhenMLXStarted = espressoIsLoaded + return "mlx output" + } + ) + + XCTAssertEqual(result.value, "mlx output") + XCTAssertTrue(result.usedMLX) + XCTAssertFalse(espressoIsLoaded) + XCTAssertFalse(espressoWasLoadedWhenMLXStarted) + } + func testDisabledFallbackDoesNotRunMLX() async { + var preparedForFallback = false var ranMLX = false do { _ = try await TextProcessor.runEspressoWithMLXFallback( fallbackEnabled: false, espresso: { throw StubError.espresso }, + prepareForMLXFallback: { + preparedForFallback = true + }, mlx: { ranMLX = true return "mlx output" @@ -71,6 +96,7 @@ final class EspressoFallbackTests: XCTestCase { XCTFail("Expected the Espresso failure") } catch { XCTAssertEqual(error.localizedDescription, "espresso failed") + XCTAssertFalse(preparedForFallback) XCTAssertFalse(ranMLX) } } @@ -214,11 +240,13 @@ final class EspressoFallbackTests: XCTestCase { temperature: 0 ) if index == 0 { + outcome = await processor.consumeEspressoOutcome() + let espressoIsLoaded = await processor.espressoLLM.isLoaded + XCTAssertFalse(espressoIsLoaded) baselineFootprint = currentMemoryFootprint() options.localLLMBackend = .mlx } } - outcome = await processor.consumeEspressoOutcome() } XCTAssertFalse(output.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty)