diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift index 739d0944..3b1618c3 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift @@ -202,6 +202,15 @@ extension InferenceEngine { public var supportsLogits: Bool { false } } +extension InferenceEngine { + /// Guided/structured generation needs either per-step logits (CPU-side + /// constrained decoding) or GPU-side constrained sampling + /// (`ConstrainedGenerationCapable`). + public var supportsGuidedGeneration: Bool { + supportsLogits || self is any ConstrainedGenerationCapable + } +} + extension InferenceEngine { /// Default: no prefix hits (engine doesn't track history). public var lastPrefixHitCount: Int { 0 } diff --git a/swift/Sources/CoreAILanguageModels/LanguageModel/CoreAILanguageModel.swift b/swift/Sources/CoreAILanguageModels/LanguageModel/CoreAILanguageModel.swift index 95d3fd28..3a8a8174 100644 --- a/swift/Sources/CoreAILanguageModels/LanguageModel/CoreAILanguageModel.swift +++ b/swift/Sources/CoreAILanguageModels/LanguageModel/CoreAILanguageModel.swift @@ -199,15 +199,10 @@ public struct CoreAILanguageModel: LanguageModel { resources.unloadResources() } - /// Whether guided generation is available for this model. + /// Whether guided generation is available for this model. Assumes yes + /// before an engine is loaded. private var isGuidedGenerationSupported: Bool { - if let isConstrainedCapable = resources.loadedEngineIsConstrainedCapable { - return isConstrainedCapable - } - if let supportsLogits = resources.loadedEngineSupportsLogits { - return supportsLogits - } - return true + resources.loadedEngineSupportsGuidedGeneration ?? true } // MARK: - Executor @@ -283,7 +278,7 @@ public struct CoreAILanguageModel: LanguageModel { // Check if guided generation is requested if let schema = request.schema { - guard engine.supportsLogits || engine is any ConstrainedGenerationCapable else { + guard engine.supportsGuidedGeneration else { throw LanguageModelError.unsupportedCapability( .init( capability: .guidedGeneration, diff --git a/swift/Sources/CoreAILanguageModels/LanguageModel/ModelResources.swift b/swift/Sources/CoreAILanguageModels/LanguageModel/ModelResources.swift index bc050652..7d5c4aa3 100644 --- a/swift/Sources/CoreAILanguageModels/LanguageModel/ModelResources.swift +++ b/swift/Sources/CoreAILanguageModels/LanguageModel/ModelResources.swift @@ -36,18 +36,10 @@ final class ModelResources: ResourceManaging { var isLoaded: Bool { state.withLock { $0.loaded != nil } } - /// `supportsLogits` of the resident engine, or `nil` when nothing is loaded. - /// Used only for best-effort capability reporting before a load. - var loadedEngineSupportsLogits: Bool? { - state.withLock { $0.loaded?.supportsLogits } - } - - /// Whether the loaded engine supports GPU-side constrained generation, or `nil` when unloaded. - var loadedEngineIsConstrainedCapable: Bool? { - state.withLock { engine in - guard let loaded = engine.loaded else { return nil } - return loaded is any ConstrainedGenerationCapable - } + /// Guided-generation capability of the resident engine, or `nil` when nothing + /// is loaded. Mirrors the check the executor enforces per-request. + var loadedEngineSupportsGuidedGeneration: Bool? { + state.withLock { $0.loaded?.supportsGuidedGeneration } } /// Returns the engine, loading it on first use. Concurrent callers share one diff --git a/swift/Tests/LanguageModelsTests/ModelResourcesTests.swift b/swift/Tests/LanguageModelsTests/ModelResourcesTests.swift index 334390c5..02e00801 100644 --- a/swift/Tests/LanguageModelsTests/ModelResourcesTests.swift +++ b/swift/Tests/LanguageModelsTests/ModelResourcesTests.swift @@ -107,6 +107,42 @@ struct ModelResourcesTests { #expect(resources.isLoaded) } + @Test( + "loadedEngineSupportsGuidedGeneration stays true after loading a logits-only engine" + ) + func guidedGenerationSupportedForLogitsOnlyEngineAfterLoad() async throws { + // MockEngine(vocabSize: 100) mirrors StaticShapeEngine: supportsLogits == true, + // but it does not conform to ConstrainedGenerationCapable. + let resources = ModelResources { MockEngine(vocabSize: 100) } + + // Nothing loaded yet — best-effort reporting defaults to "supported". + #expect(resources.loadedEngineSupportsGuidedGeneration == nil) + + _ = try await resources.engine() + + // Regression check: before the fix this incorrectly returned `false` once + // an engine was loaded, because `loadedEngineIsConstrainedCapable` (false) + // shadowed `loadedEngineSupportsLogits` (true). + #expect(resources.loadedEngineSupportsGuidedGeneration == true) + } + + @Test( + "loadedEngineSupportsGuidedGeneration is true for a constrained-capable engine without logits" + ) + func guidedGenerationSupportedForConstrainedOnlyEngineAfterLoad() async throws { + // MockConstrainedEngine conforms to ConstrainedGenerationCapable but has + // supportsLogits == false — the other branch of the capability OR. Guided + // generation must still be reported supported via GPU-side constrained + // sampling. + let resources = ModelResources { MockConstrainedEngine(scriptedTokens: []) } + + #expect(resources.loadedEngineSupportsGuidedGeneration == nil) + + _ = try await resources.engine() + + #expect(resources.loadedEngineSupportsGuidedGeneration == true) + } + @Test("unloadResources during an active borrow defers teardown until it finishes") func unloadDeferredDuringActiveBorrow() async throws { let resources = ModelResources { MockEngine() }