From bc06f4fcc8604981e0a101c7f6b659d9ba53d067 Mon Sep 17 00:00:00 2001 From: john-rocky Date: Tue, 7 Jul 2026 23:32:03 +0900 Subject: [PATCH 1/6] Add LiteRT trait and LiteRTLanguageModel backend Runs .litertlm models (e.g. Gemma 4) fully on-device via Google's LiteRT-LM runtime, with Metal GPU acceleration on iOS and macOS. - LiteRT package trait, gating a swift-litert-lm dependency to iOS/macOS; default builds are unaffected - LiteRTLanguageModel: respond/streamResponse, image inputs for models with a vision tower, structured generation via schema-in-prompt + JSON extraction, and prompt-driven tool calling for respond (with the ToolExecutionDecision delegate flow) - Env-gated tests (LITERT_TEST_MODEL), README section and provider table updates --- Package.resolved | 11 +- Package.swift | 7 + README.md | 38 +- .../Models/LiteRTLanguageModel.swift | 672 ++++++++++++++++++ .../LiteRTLanguageModelTests.swift | 63 ++ 5 files changed, 789 insertions(+), 2 deletions(-) create mode 100644 Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift create mode 100644 Tests/AnyLanguageModelTests/LiteRTLanguageModelTests.swift diff --git a/Package.resolved b/Package.resolved index b689a691..1d450a68 100644 --- a/Package.resolved +++ b/Package.resolved @@ -1,5 +1,5 @@ { - "originHash" : "8c8059604845311c40f47c771a21887b94c4a9dffc50f0ffdeceef9ef2f3c866", + "originHash" : "6173d2b219d89622a95aae2fcb7d8bde21639f4c400f469b7ec9877bfc9f2987", "pins" : [ { "identity" : "eventsource", @@ -46,6 +46,15 @@ "version" : "1.6.0" } }, + { + "identity" : "swift-litert-lm", + "kind" : "remoteSourceControl", + "location" : "https://github.com/john-rocky/swift-litert-lm", + "state" : { + "revision" : "50c524e925c24e2651aaad1ae6baee82a15ab6b3", + "version" : "0.1.1" + } + }, { "identity" : "swift-nio", "kind" : "remoteSourceControl", diff --git a/Package.swift b/Package.swift index b559f0cf..8be34477 100644 --- a/Package.swift +++ b/Package.swift @@ -25,6 +25,7 @@ let package = Package( .trait(name: "CoreML"), .trait(name: "MLX"), .trait(name: "Llama"), + .trait(name: "LiteRT"), .trait(name: "AsyncHTTPClient"), .default(enabledTraits: []), ], @@ -43,6 +44,7 @@ let package = Package( .package(url: "https://github.com/mattt/llama.swift", .upToNextMajor(from: "2.7484.0")), .package(url: "https://github.com/mattt/PartialJSONDecoder", from: "1.0.0"), .package(url: "https://github.com/ml-explore/mlx-swift-lm", from: "3.0.0"), + .package(url: "https://github.com/john-rocky/swift-litert-lm", from: "0.1.1"), .package(url: "https://github.com/swiftlang/swift-syntax", from: "602.0.0"), .package(url: "https://github.com/swift-server/async-http-client.git", from: "1.24.0"), ], @@ -94,6 +96,11 @@ let package = Package( package: "llama.swift", condition: .when(traits: ["Llama"]) ), + .product( + name: "LiteRTFoundation", + package: "swift-litert-lm", + condition: .when(platforms: [.iOS, .macOS], traits: ["LiteRT"]) + ), .product( name: "AsyncHTTPClient", package: "async-http-client", diff --git a/README.md b/README.md index 41e75842..b53dec1b 100644 --- a/README.md +++ b/README.md @@ -73,6 +73,7 @@ session.toolExecutionDelegate = ToolExecutionObserver() - [x] [Core ML](https://developer.apple.com/documentation/coreml) models - [x] [MLX](https://github.com/ml-explore/mlx-swift) models - [x] [llama.cpp](https://github.com/ggml-org/llama.cpp) (GGUF models) +- [x] [LiteRT-LM](https://github.com/google-ai-edge/litert-lm) (`.litertlm` models, via [swift-litert-lm](https://github.com/john-rocky/swift-litert-lm)) - [x] Ollama [HTTP API](https://github.com/ollama/ollama/blob/main/docs/api.md) - [x] Anthropic [Messages API](https://docs.claude.com/en/api/messages) - [x] Google [Gemini API](https://ai.google.dev/api/generate-content) @@ -117,6 +118,8 @@ This results in smaller binary sizes and faster build times. (depends on `ml-explore/mlx-swift-lm`) - `Llama`: Enables llama.cpp support (requires `mattt/llama.swift`) +- `LiteRT`: Enables LiteRT-LM support for Gemma 4 and other `.litertlm` models + (requires `john-rocky/swift-litert-lm`; iOS and macOS only) By default, no traits are enabled. To enable specific traits, specify them in your package's dependencies: @@ -366,13 +369,14 @@ Image support varies by provider: | Core ML | — | | MLX | model-dependent | | llama.cpp | — | +| LiteRT-LM | model-dependent | | Ollama | model-dependent | | OpenAI | yes | | Open Responses | yes | | Anthropic | yes | | Google Gemini | yes | -For MLX and Ollama, +For MLX, LiteRT-LM, and Ollama, use a vision-capable model (for example, a VLM or `-vl` variant). @@ -610,6 +614,38 @@ let response = try await session.respond( ) ``` +### LiteRT-LM + +Runs `.litertlm` models (for example, Gemma 4) fully on-device via Google's +[LiteRT-LM](https://github.com/google-ai-edge/litert-lm) runtime with Metal GPU +acceleration (requires `LiteRT` trait; iOS and macOS only): + +```swift +let model = LiteRTLanguageModel(model: .gemma4_E2B) + +let session = LanguageModelSession(model: model) +let response = try await session.respond(to: "What is the capital of France?") +``` + +The `.litertlm` file is downloaded from Hugging Face on first use and cached +under Application Support. You can also load any Hugging Face repo hosting a +`.litertlm`, or a local file: + +```swift +// Any Hugging Face repo +let model = LiteRTLanguageModel( + huggingFaceRepo: "litert-community/gemma-4-E4B-it-litert-lm", + fileName: "gemma-4-E4B-it.litertlm") + +// Local file (for example, a fine-tuned model) +let model = LiteRTLanguageModel(modelFileURL: url) +``` + +Image inputs are supported for models that ship a vision tower (pass +`modalities: .textImage`). Structured generation is prompt-driven — the JSON +schema is included in the prompt and the response is parsed. Tool calling is +supported for `respond` (not yet for streaming). + ### Ollama Run models locally via Ollama's diff --git a/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift b/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift new file mode 100644 index 00000000..22e60aab --- /dev/null +++ b/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift @@ -0,0 +1,672 @@ +import Foundation + +#if LiteRT + @preconcurrency import LiteRTFoundation + + /// A language model that runs `.litertlm` models fully on-device via Google's + /// [LiteRT-LM](https://github.com/google-ai-edge/litert-lm) runtime. + /// + /// Use this model to run Gemma 4 (and other LiteRT-LM models) on iOS and macOS + /// with Metal GPU acceleration, including image understanding for models that + /// ship a vision tower. + /// + /// ```swift + /// let model = LiteRTLanguageModel(model: .gemma4_E2B) + /// let session = LanguageModelSession(model: model) + /// let response = try await session.respond(to: "What is the capital of France?") + /// ``` + /// + /// The model file is downloaded from Hugging Face on first use and cached under + /// Application Support. Loading is lazy: the engine is brought up on the first + /// request (or when ``prewarm(for:promptPrefix:)`` is called). + /// + /// Structured generation is prompt-driven (the JSON schema is included in the + /// prompt and the response is parsed), and tool calling is supported for + /// ``respond(within:to:generating:includeSchemaInPrompt:options:)`` + /// (not yet for streaming). + public struct LiteRTLanguageModel: LanguageModel { + /// The reason the model is unavailable. + /// This model is always available; loading errors surface when responding. + public typealias UnavailableReason = Never + + private let engine: LazyEngine + + /// Creates a model from the swift-litert-lm catalog, downloading the + /// `.litertlm` from Hugging Face on first use. + /// + /// - Parameters: + /// - model: The catalog model to run (for example, `.gemma4_E2B`). + /// - modalities: Towers to enable. Defaults to the model's default + /// modalities. Requesting an unsupported tower is ignored. + /// - storageDirectory: Where to keep the downloaded model. Defaults to + /// Application Support/LiteRTModels. + /// - allowUnsafeMemory: Bypass the device-RAM safety check. + /// - maxTokens: Context (KV cache) budget. Defaults to the catalog value. + /// - onDownloadProgress: Called during the first-run model download. + public init( + model: LiteRTModel, + modalities: Modality? = nil, + storageDirectory: URL? = nil, + allowUnsafeMemory: Bool = false, + maxTokens: Int? = nil, + onDownloadProgress: (@Sendable (ModelDownloader.Progress) -> Void)? = nil + ) { + self.engine = LazyEngine { + if !allowUnsafeMemory { + let ram = Int64(ProcessInfo.processInfo.physicalMemory) + if ram < model.minimumDeviceRAM { + throw LiteRTChatError.insufficientMemory( + haveBytes: ram, + needBytes: model.minimumDeviceRAM + ) + } + } + let path = try await LiteRTChat.ensureModel( + model, + storageDirectory: storageDirectory, + onProgress: onDownloadProgress + ) + var wanted = modalities ?? model.defaultModalities + wanted.formIntersection(model.supportedModalities) + return try await makeEngine( + modelPath: path, + modalities: wanted, + visionBackend: model.visionBackend, + audioBackend: model.audioBackend, + maxTokens: maxTokens ?? model.defaultMaxTokens, + visualTokenBudget: model.defaultVisualTokenBudget + ) + } + } + + /// Creates a model from a local `.litertlm` file. No download. + /// + /// - Parameters: + /// - modelFileURL: Absolute file URL of an on-disk `.litertlm`. + /// - modalities: Towers to bring up. Defaults to `.all`; only the ones + /// the model actually contains will work. + /// - visionBackend: Backend for the vision encoder. Defaults to `.cpu()` + /// (the safe choice for Gemma 4 on iOS). + /// - audioBackend: Backend for the audio encoder. Defaults to `.cpu()`. + /// - visualTokenBudget: Per-image visual-token cap (`nil` = engine default). + /// - maxTokens: Context (KV cache) budget. + public init( + modelFileURL: URL, + modalities: Modality = .all, + visionBackend: Backend = .cpu(), + audioBackend: Backend = .cpu(), + visualTokenBudget: Int32? = nil, + maxTokens: Int = 2048 + ) { + self.engine = LazyEngine { + guard FileManager.default.fileExists(atPath: modelFileURL.path) else { + throw LiteRTChatError.modelFileNotFound(modelFileURL) + } + return try await makeEngine( + modelPath: modelFileURL.path, + modalities: modalities, + visionBackend: visionBackend, + audioBackend: audioBackend, + maxTokens: maxTokens, + visualTokenBudget: visualTokenBudget + ) + } + } + + /// Creates a model from any Hugging Face repo hosting a `.litertlm`, + /// downloading it on first use. + /// + /// - Parameters: + /// - huggingFaceRepo: For example, `"litert-community/gemma-4-E4B-it-litert-lm"`. + /// - fileName: The `.litertlm` file in that repo. + /// - revision: Git revision or branch. Defaults to `main`. + /// - modalities: Defaults to text-only (`[]`) — the safe choice for an + /// unknown model. Pass `.textImage` or `.all` if the model ships those + /// encoders. + /// - visionBackend: Backend for the vision encoder. Defaults to `.cpu()`. + /// - audioBackend: Backend for the audio encoder. Defaults to `.cpu()`. + /// - visualTokenBudget: Per-image visual-token cap (`nil` = engine default). + /// - maxTokens: Context (KV cache) budget. + /// - storageDirectory: Where to keep the downloaded model. + /// - onDownloadProgress: Called during the first-run model download. + public init( + huggingFaceRepo: String, + fileName: String, + revision: String = "main", + modalities: Modality = [], + visionBackend: Backend = .cpu(), + audioBackend: Backend = .cpu(), + visualTokenBudget: Int32? = nil, + maxTokens: Int = 2048, + storageDirectory: URL? = nil, + onDownloadProgress: (@Sendable (ModelDownloader.Progress) -> Void)? = nil + ) { + self.engine = LazyEngine { + guard + let url = URL( + string: + "https://huggingface.co/\(huggingFaceRepo)/resolve/\(revision)/\(fileName)?download=true" + ) + else { + throw LiteRTChatError.modelFileNotFound(URL(fileURLWithPath: fileName)) + } + let directory = try storageDirectory ?? LiteRTChat.defaultStorageDirectory() + let destination = directory.appendingPathComponent(fileName) + try await ModelDownloader.shared.download( + from: url, + to: destination, + expectedBytes: nil, + onProgress: onDownloadProgress + ) + return try await makeEngine( + modelPath: destination.path, + modalities: modalities, + visionBackend: visionBackend, + audioBackend: audioBackend, + maxTokens: maxTokens, + visualTokenBudget: visualTokenBudget + ) + } + } + + public func prewarm( + for session: LanguageModelSession, + promptPrefix: Prompt? + ) { + let engine = self.engine + Task { _ = try? await engine.ready() } + } + + public func respond( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) async throws -> LanguageModelSession.Response where Content: Generable { + let engine = try await self.engine.ready() + + let schemaJSON: String? + if type == String.self { + schemaJSON = nil + } else { + schemaJSON = try encodeSchema(type.generationSchema) + } + + let tools = session.tools + var plan = makePlan( + from: session.transcript, + fallbackPrompt: prompt.description, + schemaJSON: includeSchemaInPrompt ? schemaJSON : nil, + tools: tools + ) + let sampler = makeSampler(for: options, structured: schemaJSON != nil || !tools.isEmpty) + + var entries: [Transcript.Entry] = [] + var text = "" + var toolRounds = 0 + + while true { + let conversation = try await engine.createConversation( + with: ConversationConfig( + systemMessage: plan.systemMessage, + initialMessages: plan.history, + samplerConfig: sampler + ) + ) + + text = "" + for try await chunk in conversation.sendMessageStream(plan.prompt) { + text += chunk.toString + } + + guard !tools.isEmpty, + toolRounds < maxToolRounds, + let parsed = parseToolCall(from: text, tools: tools) + else { break } + toolRounds += 1 + + let resolution = try await resolveToolCall( + name: parsed.name, + argumentsJSON: parsed.arguments, + session: session + ) + switch resolution { + case .stop(let call): + entries.append(.toolCalls(Transcript.ToolCalls([call]))) + return LanguageModelSession.Response( + content: "" as! Content, + rawContent: GeneratedContent(""), + transcriptEntries: ArraySlice(entries) + ) + case .invocation(let call, let output): + entries.append(.toolCalls(Transcript.ToolCalls([call]))) + entries.append(.toolOutput(output)) + plan = plan.continuing(afterModelText: text, toolOutput: output) + } + } + + if type == String.self { + return LanguageModelSession.Response( + content: text as! Content, + rawContent: GeneratedContent(text), + transcriptEntries: ArraySlice(entries) + ) + } + + let json = extractJSONObject(from: text) ?? text + let generatedContent = try GeneratedContent(json: json) + let content = try type.init(generatedContent) + return LanguageModelSession.Response( + content: content, + rawContent: generatedContent, + transcriptEntries: ArraySlice(entries) + ) + } + + public func streamResponse( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) -> sending LanguageModelSession.ResponseStream where Content: Generable { + let lazyEngine = self.engine + let stream: AsyncThrowingStream.Snapshot, any Error> = + AsyncThrowingStream { continuation in + let task = Task { + do { + let engine = try await lazyEngine.ready() + + let schemaJSON: String? + if type == String.self { + schemaJSON = nil + } else { + schemaJSON = try encodeSchema(type.generationSchema) + } + + let plan = makePlan( + from: session.transcript, + fallbackPrompt: prompt.description, + schemaJSON: includeSchemaInPrompt ? schemaJSON : nil, + tools: [] + ) + let conversation = try await engine.createConversation( + with: ConversationConfig( + systemMessage: plan.systemMessage, + initialMessages: plan.history, + samplerConfig: makeSampler(for: options, structured: schemaJSON != nil) + ) + ) + + var text = "" + for try await chunk in conversation.sendMessageStream(plan.prompt) { + let delta = chunk.toString + guard !delta.isEmpty else { continue } + text += delta + + if type == String.self { + continuation.yield( + .init( + content: (text as! Content).asPartiallyGenerated(), + rawContent: GeneratedContent(text) + ) + ) + } else if let json = extractJSONObject(from: text), + let raw = try? GeneratedContent(json: json), + let parsed = try? type.init(raw) + { + continuation.yield( + .init( + content: parsed.asPartiallyGenerated(), + rawContent: raw + ) + ) + } else { + // Structured responses stream as incomplete JSON fragments. + // Skip snapshots until the accumulated JSON parses cleanly. + } + } + + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } + + continuation.onTermination = { _ in + task.cancel() + } + } + + return LanguageModelSession.ResponseStream(stream: stream) + } + } + + // MARK: - Engine Bring-Up + + /// Brings up the engine on first use and shares it across requests. + /// Engine bring-up loads multi-GB weights, so it must happen exactly once. + private actor LazyEngine { + private var task: Task? + private let bringUp: @Sendable () async throws -> Engine + + init(_ bringUp: @escaping @Sendable () async throws -> Engine) { + self.bringUp = bringUp + } + + func ready() async throws -> Engine { + if task == nil { + let bringUp = self.bringUp + task = Task { try await bringUp() } + } + return try await task!.value + } + } + + private func makeEngine( + modelPath: String, + modalities: Modality, + visionBackend: Backend, + audioBackend: Backend, + maxTokens: Int, + visualTokenBudget: Int32? + ) async throws -> Engine { + ExperimentalFlags.optIntoExperimentalAPIs() + if modalities.contains(.vision), let visualTokenBudget { + ExperimentalFlags.visualTokenBudget = visualTokenBudget + } + let caches = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask).first + let config = try EngineConfig( + modelPath: modelPath, + backend: .gpu, + visionBackend: modalities.contains(.vision) ? visionBackend : nil, + audioBackend: modalities.contains(.audio) ? audioBackend : nil, + maxNumTokens: maxTokens, + cacheDir: caches?.path, + // The engine default is 1 image per conversation (a 2nd image + // overwrites the 1st); allow several so multi-image chats work. + maxNumImages: modalities.contains(.vision) ? 16 : nil + ) + let engine = Engine(engineConfig: config) + try await engine.initialize() + return engine + } + + // MARK: - Transcript → LiteRT Messages + + private struct GenerationPlan { + var systemMessage: Message? + var history: [Message] + var prompt: Message + + /// Extends the plan after a tool round-trip: the trigger prompt and the + /// model's tool-call text become history, and the tool result becomes the + /// new trigger. + func continuing(afterModelText text: String, toolOutput: Transcript.ToolOutput) -> GenerationPlan { + var history = self.history + history.append(prompt) + history.append(Message(text, role: .model)) + let result = textContent(of: toolOutput.segments) + let trigger = Message( + "Tool \"\(toolOutput.toolName)\" returned: \(result)\nUse this result to answer the user.", + role: .user + ) + return GenerationPlan(systemMessage: systemMessage, history: history, prompt: trigger) + } + } + + /// Splits the session transcript into a system message, prior turns, and the + /// message to generate from. The generation trigger is the last `.prompt` or + /// (in a tool round-trip) the last `.toolOutput` entry. + private func makePlan( + from transcript: Transcript, + fallbackPrompt: String, + schemaJSON: String?, + tools: [any Tool] + ) -> GenerationPlan { + let entries = Array(transcript) + let triggerIndex = entries.lastIndex { entry in + switch entry { + case .prompt, .toolOutput: return true + default: return false + } + } + + var systemText: [String] = [] + if !tools.isEmpty { + systemText.append(toolInstructions(tools)) + } + var history: [Message] = [] + var trigger: Message? + + for (index, entry) in entries.enumerated() { + let isTrigger = (index == triggerIndex) + switch entry { + case .instructions(let instructions): + systemText.append(textContent(of: instructions.segments)) + case .prompt(let prompt): + var contents = messageContents(of: prompt.segments) + if isTrigger, let schemaJSON, !schemaJSON.isEmpty { + contents.append( + .text( + "\n\nRespond with ONLY a JSON object that conforms to this JSON schema. " + + "Output valid JSON and nothing else:\n\(schemaJSON)" + ) + ) + } + let message = Message(contents: contents, role: .user) + if isTrigger { trigger = message } else { history.append(message) } + case .response(let response): + history.append(Message(contents: [.text(textContent(of: response.segments))], role: .model)) + case .toolOutput(let output): + let result = textContent(of: output.segments) + let message = Message( + "Tool \"\(output.toolName)\" returned: \(result)\nUse this result to answer the user.", + role: .user + ) + if isTrigger { trigger = message } else { history.append(message) } + case .toolCalls: + history.append(Message("[the assistant called a tool]", role: .model)) + } + } + + let system = systemText.joined(separator: "\n").trimmingCharacters(in: .whitespacesAndNewlines) + return GenerationPlan( + systemMessage: system.isEmpty ? nil : Message(system, role: .system), + history: history, + prompt: trigger ?? Message(fallbackPrompt, role: .user) + ) + } + + /// Maps transcript segments to LiteRT content: text, structured content as + /// JSON text, and images. + private func messageContents(of segments: [Transcript.Segment]) -> [Content] { + var contents: [Content] = [] + for segment in segments { + switch segment { + case .text(let text): + if !text.content.isEmpty { + contents.append(.text(text.content)) + } + case .structure(let structure): + contents.append(.text(structure.content.jsonString)) + case .image(let image): + switch image.source { + case .data(let data, _): + contents.append(.imageData(data)) + case .url(let url): + if url.isFileURL { + contents.append(.imageFile(url.path)) + } else if let data = try? Data(contentsOf: url) { + contents.append(.imageData(data)) + } + } + } + } + return contents.isEmpty ? [.text("")] : contents + } + + /// Concatenates the text of a segment list (non-text segments are ignored). + private func textContent(of segments: [Transcript.Segment]) -> String { + segments.compactMap { segment in + if case .text(let text) = segment { return text.content } else { return nil } + }.joined(separator: " ") + } + + // MARK: - Structured Generation + + private func encodeSchema(_ schema: GenerationSchema) throws -> String { + let resolvedSchema = schema.withResolvedRoot() ?? schema + let data = try JSONEncoder().encode(resolvedSchema) + return String(data: data, encoding: .utf8) ?? "" + } + + /// Extracts the first balanced JSON object from model text + /// (strips surrounding prose and code fences). + private func extractJSONObject(from text: String) -> String? { + guard let start = text.firstIndex(of: "{") else { return nil } + var depth = 0 + var inString = false + var escaped = false + var index = start + while index < text.endIndex { + let character = text[index] + if inString { + if escaped { + escaped = false + } else if character == "\\" { + escaped = true + } else if character == "\"" { + inString = false + } + } else if character == "\"" { + inString = true + } else if character == "{" { + depth += 1 + } else if character == "}" { + depth -= 1 + if depth == 0 { return String(text[start ... index]) } + } + index = text.index(after: index) + } + return nil + } + + // MARK: - Sampling + + private func makeSampler(for options: GenerationOptions, structured: Bool) -> SamplerConfig? { + var topK = 40 + var topP = 0.95 + // Lower default temperature for structured / tool generation + // (more reliable JSON). + var temperature = structured ? 0.0 : 0.8 + if let explicit = options.temperature { + temperature = explicit + } + if let sampling = options.sampling { + switch sampling.mode { + case .greedy: + temperature = 0.0 + case .topK(let k, _): + topK = k + case .nucleus(let probabilityThreshold, _): + topP = probabilityThreshold + } + } + return try? SamplerConfig(topK: topK, topP: Float(topP), temperature: Float(temperature)) + } + + // MARK: - Tool Calling + + private let maxToolRounds = 4 + + /// Describes the enabled tools and the tool-call JSON format for the prompt. + private func toolInstructions(_ tools: [any Tool]) -> String { + var lines = ["You can call tools to help answer the user. Available tools:"] + for tool in tools { + let parameters = (try? encodeSchema(tool.parameters)) ?? "{}" + lines.append("- \(tool.name): \(tool.description). arguments schema: \(parameters)") + } + lines.append( + "To call a tool, reply with ONLY this JSON and nothing else: " + + "{\"tool_call\": {\"name\": \"\", \"arguments\": { ... }}}. " + + "If no tool is needed, answer the user directly." + ) + return lines.joined(separator: "\n") + } + + /// Parses a tool call from model output, if present and naming a known tool. + private func parseToolCall( + from text: String, + tools: [any Tool] + ) -> (name: String, arguments: String)? { + guard let json = extractJSONObject(from: text), + let data = json.data(using: .utf8), + let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let call = object["tool_call"] as? [String: Any], + let name = call["name"] as? String, + tools.contains(where: { $0.name == name }) + else { return nil } + let arguments = call["arguments"] ?? [String: Any]() + let argumentsData = (try? JSONSerialization.data(withJSONObject: arguments)) ?? Data("{}".utf8) + return (name, String(data: argumentsData, encoding: .utf8) ?? "{}") + } + + private enum ToolResolution { + case stop(call: Transcript.ToolCall) + case invocation(call: Transcript.ToolCall, output: Transcript.ToolOutput) + } + + private func resolveToolCall( + name: String, + argumentsJSON: String, + session: LanguageModelSession + ) async throws -> ToolResolution { + let arguments = (try? GeneratedContent(json: argumentsJSON)) ?? GeneratedContent(properties: [:]) + let call = Transcript.ToolCall(id: UUID().uuidString, toolName: name, arguments: arguments) + + if let delegate = session.toolExecutionDelegate { + await delegate.didGenerateToolCalls([call], in: session) + } + + var decision: ToolExecutionDecision = .execute + if let delegate = session.toolExecutionDelegate { + decision = await delegate.toolCallDecision(for: call, in: session) + } + + switch decision { + case .stop: + return .stop(call: call) + case .provideOutput(let segments): + let output = Transcript.ToolOutput(id: call.id, toolName: call.toolName, segments: segments) + if let delegate = session.toolExecutionDelegate { + await delegate.didExecuteToolCall(call, output: output, in: session) + } + return .invocation(call: call, output: output) + case .execute: + guard let tool = session.tools.first(where: { $0.name == name }) else { + let message = Transcript.Segment.text(.init(content: "Tool not found: \(name)")) + let output = Transcript.ToolOutput(id: call.id, toolName: name, segments: [message]) + if let delegate = session.toolExecutionDelegate { + await delegate.didExecuteToolCall(call, output: output, in: session) + } + return .invocation(call: call, output: output) + } + + do { + let segments = try await tool.makeOutputSegments(from: call.arguments) + let output = Transcript.ToolOutput(id: call.id, toolName: tool.name, segments: segments) + if let delegate = session.toolExecutionDelegate { + await delegate.didExecuteToolCall(call, output: output, in: session) + } + return .invocation(call: call, output: output) + } catch { + if let delegate = session.toolExecutionDelegate { + await delegate.didFailToolCall(call, error: error, in: session) + } + throw LanguageModelSession.ToolCallError(tool: tool, underlyingError: error) + } + } + } +#endif diff --git a/Tests/AnyLanguageModelTests/LiteRTLanguageModelTests.swift b/Tests/AnyLanguageModelTests/LiteRTLanguageModelTests.swift new file mode 100644 index 00000000..dd9a92a5 --- /dev/null +++ b/Tests/AnyLanguageModelTests/LiteRTLanguageModelTests.swift @@ -0,0 +1,63 @@ +import Foundation +import Testing + +@testable import AnyLanguageModel + +#if LiteRT + /// Path to a local `.litertlm` file to test against + /// (for example, gemma-4-E2B-it.litertlm). + /// These tests load multi-GB weights, so they only run when explicitly + /// requested via the `LITERT_TEST_MODEL` environment variable. + private let liteRTTestModelPath = ProcessInfo.processInfo.environment["LITERT_TEST_MODEL"] + + private let shouldRunLiteRTTests = liteRTTestModelPath != nil + + @Generable + private struct CityAnswer { + var city: String + } + + @Suite("LiteRTLanguageModel", .enabled(if: shouldRunLiteRTTests), .serialized) + struct LiteRTLanguageModelTests { + private var model: LiteRTLanguageModel { + LiteRTLanguageModel( + modelFileURL: URL(fileURLWithPath: liteRTTestModelPath!), + modalities: [] + ) + } + + @Test func respondToTextPrompt() async throws { + let session = LanguageModelSession(model: model) + let response = try await session.respond(to: "Reply with a single word: hello") + #expect(!response.content.isEmpty) + } + + @Test func streamTextPrompt() async throws { + let session = LanguageModelSession(model: model) + var snapshots = 0 + var lastContent = "" + for try await snapshot in session.streamResponse(to: "Count from 1 to 5.") { + snapshots += 1 + lastContent = snapshot.content + } + #expect(snapshots > 1) + #expect(!lastContent.isEmpty) + } + + @Test func structuredGeneration() async throws { + let session = LanguageModelSession(model: model) + let response = try await session.respond( + to: "What is the capital of France?", + generating: CityAnswer.self + ) + #expect(!response.content.city.isEmpty) + } + + @Test func multiTurnConversationKeepsContext() async throws { + let session = LanguageModelSession(model: model) + _ = try await session.respond(to: "My name is Alice. Remember it.") + let response = try await session.respond(to: "What is my name?") + #expect(response.content.localizedCaseInsensitiveContains("alice")) + } + } +#endif From bfbfd034d97206fd49297cdea8ce6c730ea4e254 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Fri, 11 Sep 2026 04:26:10 -0700 Subject: [PATCH 2/6] Use the official LiteRT-LM Swift package --- Package.resolved | 37 +++- Package.swift | 8 +- README.md | 16 +- .../Models/LiteRTLanguageModel.swift | 174 ++++++------------ .../LiteRTLanguageModelTests.swift | 38 +++- 5 files changed, 130 insertions(+), 143 deletions(-) diff --git a/Package.resolved b/Package.resolved index 1d450a68..75dfacc9 100644 --- a/Package.resolved +++ b/Package.resolved @@ -1,5 +1,5 @@ { - "originHash" : "6173d2b219d89622a95aae2fcb7d8bde21639f4c400f469b7ec9877bfc9f2987", + "originHash" : "802c7ad583bfa4b4de0e754641d0e364438151ba6bdd8bcbcd9050572a0cfc4f", "pins" : [ { "identity" : "eventsource", @@ -19,6 +19,15 @@ "version" : "1.3.1" } }, + { + "identity" : "litert-lm", + "kind" : "remoteSourceControl", + "location" : "https://github.com/google-ai-edge/LiteRT-LM", + "state" : { + "revision" : "e9fd8c53ff968071774206163027dd84bedfe925", + "version" : "0.17.0" + } + }, { "identity" : "partialjsondecoder", "kind" : "remoteSourceControl", @@ -28,6 +37,15 @@ "version" : "1.0.0" } }, + { + "identity" : "swift-asn1", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-asn1.git", + "state" : { + "revision" : "d9a5b37470adc940d22c3bcd5ca6953a516b727f", + "version" : "1.7.2" + } + }, { "identity" : "swift-atomics", "kind" : "remoteSourceControl", @@ -47,12 +65,21 @@ } }, { - "identity" : "swift-litert-lm", + "identity" : "swift-crypto", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-crypto.git", + "state" : { + "revision" : "da9d28d69ebe3894b18376c8f2395c2f37b8448f", + "version" : "4.5.2" + } + }, + { + "identity" : "swift-huggingface", "kind" : "remoteSourceControl", - "location" : "https://github.com/john-rocky/swift-litert-lm", + "location" : "https://github.com/huggingface/swift-huggingface", "state" : { - "revision" : "50c524e925c24e2651aaad1ae6baee82a15ab6b3", - "version" : "0.1.1" + "revision" : "b5403ed09403f674601fd1123e07c5b32914d16f", + "version" : "0.10.1" } }, { diff --git a/Package.swift b/Package.swift index 8be34477..a4f8f49c 100644 --- a/Package.swift +++ b/Package.swift @@ -44,7 +44,7 @@ let package = Package( .package(url: "https://github.com/mattt/llama.swift", .upToNextMajor(from: "2.7484.0")), .package(url: "https://github.com/mattt/PartialJSONDecoder", from: "1.0.0"), .package(url: "https://github.com/ml-explore/mlx-swift-lm", from: "3.0.0"), - .package(url: "https://github.com/john-rocky/swift-litert-lm", from: "0.1.1"), + .package(url: "https://github.com/google-ai-edge/LiteRT-LM", from: "0.17.0"), .package(url: "https://github.com/swiftlang/swift-syntax", from: "602.0.0"), .package(url: "https://github.com/swift-server/async-http-client.git", from: "1.24.0"), ], @@ -79,7 +79,7 @@ let package = Package( .product( name: "HuggingFace", package: "swift-huggingface", - condition: .when(traits: ["MLX"]) + condition: .when(traits: ["MLX", "LiteRT"]) ), .product( name: "Tokenizers", @@ -97,8 +97,8 @@ let package = Package( condition: .when(traits: ["Llama"]) ), .product( - name: "LiteRTFoundation", - package: "swift-litert-lm", + name: "LiteRTLM", + package: "LiteRT-LM", condition: .when(platforms: [.iOS, .macOS], traits: ["LiteRT"]) ), .product( diff --git a/README.md b/README.md index b53dec1b..a1527751 100644 --- a/README.md +++ b/README.md @@ -73,7 +73,7 @@ session.toolExecutionDelegate = ToolExecutionObserver() - [x] [Core ML](https://developer.apple.com/documentation/coreml) models - [x] [MLX](https://github.com/ml-explore/mlx-swift) models - [x] [llama.cpp](https://github.com/ggml-org/llama.cpp) (GGUF models) -- [x] [LiteRT-LM](https://github.com/google-ai-edge/litert-lm) (`.litertlm` models, via [swift-litert-lm](https://github.com/john-rocky/swift-litert-lm)) +- [x] [LiteRT-LM](https://github.com/google-ai-edge/litert-lm) (`.litertlm` models, using the official Swift package) - [x] Ollama [HTTP API](https://github.com/ollama/ollama/blob/main/docs/api.md) - [x] Anthropic [Messages API](https://docs.claude.com/en/api/messages) - [x] Google [Gemini API](https://ai.google.dev/api/generate-content) @@ -119,7 +119,7 @@ This results in smaller binary sizes and faster build times. - `Llama`: Enables llama.cpp support (requires `mattt/llama.swift`) - `LiteRT`: Enables LiteRT-LM support for Gemma 4 and other `.litertlm` models - (requires `john-rocky/swift-litert-lm`; iOS and macOS only) + (requires the official `google-ai-edge/LiteRT-LM` Swift package; iOS and macOS only) By default, no traits are enabled. To enable specific traits, specify them in your package's dependencies: @@ -621,28 +621,24 @@ Runs `.litertlm` models (for example, Gemma 4) fully on-device via Google's acceleration (requires `LiteRT` trait; iOS and macOS only): ```swift -let model = LiteRTLanguageModel(model: .gemma4_E2B) +let model = LiteRTLanguageModel(modelFileURL: modelURL) let session = LanguageModelSession(model: model) let response = try await session.respond(to: "What is the capital of France?") ``` -The `.litertlm` file is downloaded from Hugging Face on first use and cached -under Application Support. You can also load any Hugging Face repo hosting a -`.litertlm`, or a local file: +You can also load a `.litertlm` file from Hugging Face. The file is downloaded +on first use using the Hub client's cache and authentication: ```swift // Any Hugging Face repo let model = LiteRTLanguageModel( huggingFaceRepo: "litert-community/gemma-4-E4B-it-litert-lm", fileName: "gemma-4-E4B-it.litertlm") - -// Local file (for example, a fine-tuned model) -let model = LiteRTLanguageModel(modelFileURL: url) ``` Image inputs are supported for models that ship a vision tower (pass -`modalities: .textImage`). Structured generation is prompt-driven — the JSON +`visionBackend: .cpu()`). Structured generation is prompt-driven — the JSON schema is included in the prompt and the response is parsed. Tool calling is supported for `respond` (not yet for streaming). diff --git a/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift b/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift index 22e60aab..0ccd1c7b 100644 --- a/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift @@ -1,7 +1,9 @@ import Foundation -#if LiteRT - @preconcurrency import LiteRTFoundation +#if LiteRT && (os(iOS) || os(macOS)) && !targetEnvironment(macCatalyst) + @preconcurrency import LiteRTLM + import class HuggingFace.HubClient + import enum HuggingFace.Repo /// A language model that runs `.litertlm` models fully on-device via Google's /// [LiteRT-LM](https://github.com/google-ai-edge/litert-lm) runtime. @@ -11,13 +13,13 @@ import Foundation /// ship a vision tower. /// /// ```swift - /// let model = LiteRTLanguageModel(model: .gemma4_E2B) + /// let model = LiteRTLanguageModel(modelFileURL: modelURL) /// let session = LanguageModelSession(model: model) /// let response = try await session.respond(to: "What is the capital of France?") /// ``` /// - /// The model file is downloaded from Hugging Face on first use and cached under - /// Application Support. Loading is lazy: the engine is brought up on the first + /// Hugging Face models are downloaded on first use and stored in the Hub cache. + /// Loading is lazy: the engine is brought up on the first /// request (or when ``prewarm(for:promptPrefix:)`` is called). /// /// Structured generation is prompt-driven (the JSON schema is included in the @@ -31,140 +33,78 @@ import Foundation private let engine: LazyEngine - /// Creates a model from the swift-litert-lm catalog, downloading the - /// `.litertlm` from Hugging Face on first use. + /// Creates a model from a local `.litertlm` file. /// /// - Parameters: - /// - model: The catalog model to run (for example, `.gemma4_E2B`). - /// - modalities: Towers to enable. Defaults to the model's default - /// modalities. Requesting an unsupported tower is ignored. - /// - storageDirectory: Where to keep the downloaded model. Defaults to - /// Application Support/LiteRTModels. - /// - allowUnsafeMemory: Bypass the device-RAM safety check. - /// - maxTokens: Context (KV cache) budget. Defaults to the catalog value. - /// - onDownloadProgress: Called during the first-run model download. - public init( - model: LiteRTModel, - modalities: Modality? = nil, - storageDirectory: URL? = nil, - allowUnsafeMemory: Bool = false, - maxTokens: Int? = nil, - onDownloadProgress: (@Sendable (ModelDownloader.Progress) -> Void)? = nil - ) { - self.engine = LazyEngine { - if !allowUnsafeMemory { - let ram = Int64(ProcessInfo.processInfo.physicalMemory) - if ram < model.minimumDeviceRAM { - throw LiteRTChatError.insufficientMemory( - haveBytes: ram, - needBytes: model.minimumDeviceRAM - ) - } - } - let path = try await LiteRTChat.ensureModel( - model, - storageDirectory: storageDirectory, - onProgress: onDownloadProgress - ) - var wanted = modalities ?? model.defaultModalities - wanted.formIntersection(model.supportedModalities) - return try await makeEngine( - modelPath: path, - modalities: wanted, - visionBackend: model.visionBackend, - audioBackend: model.audioBackend, - maxTokens: maxTokens ?? model.defaultMaxTokens, - visualTokenBudget: model.defaultVisualTokenBudget - ) - } - } - - /// Creates a model from a local `.litertlm` file. No download. - /// - /// - Parameters: - /// - modelFileURL: Absolute file URL of an on-disk `.litertlm`. - /// - modalities: Towers to bring up. Defaults to `.all`; only the ones - /// the model actually contains will work. - /// - visionBackend: Backend for the vision encoder. Defaults to `.cpu()` - /// (the safe choice for Gemma 4 on iOS). - /// - audioBackend: Backend for the audio encoder. Defaults to `.cpu()`. - /// - visualTokenBudget: Per-image visual-token cap (`nil` = engine default). + /// - modelFileURL: File URL of an on-disk `.litertlm` model. + /// - backend: Backend for text generation. Defaults to Metal GPU. + /// - visionBackend: Backend for the vision encoder. Pass `.cpu()` for + /// a model with image support; `nil` disables vision. + /// - audioBackend: Backend for the audio encoder; `nil` disables audio. /// - maxTokens: Context (KV cache) budget. public init( modelFileURL: URL, - modalities: Modality = .all, - visionBackend: Backend = .cpu(), - audioBackend: Backend = .cpu(), - visualTokenBudget: Int32? = nil, + backend: Backend = .gpu, + visionBackend: Backend? = nil, + audioBackend: Backend? = nil, maxTokens: Int = 2048 ) { self.engine = LazyEngine { - guard FileManager.default.fileExists(atPath: modelFileURL.path) else { - throw LiteRTChatError.modelFileNotFound(modelFileURL) + guard modelFileURL.isFileURL, + FileManager.default.fileExists(atPath: modelFileURL.path) + else { + throw CocoaError(.fileReadNoSuchFile, userInfo: [NSURLErrorKey: modelFileURL]) } return try await makeEngine( modelPath: modelFileURL.path, - modalities: modalities, + backend: backend, visionBackend: visionBackend, audioBackend: audioBackend, - maxTokens: maxTokens, - visualTokenBudget: visualTokenBudget + maxTokens: maxTokens ) } } - /// Creates a model from any Hugging Face repo hosting a `.litertlm`, - /// downloading it on first use. + /// Creates a model from a Hugging Face repository hosting a `.litertlm` file. + /// The file is downloaded lazily using the Hub client's cache and authentication. /// /// - Parameters: - /// - huggingFaceRepo: For example, `"litert-community/gemma-4-E4B-it-litert-lm"`. - /// - fileName: The `.litertlm` file in that repo. + /// - huggingFaceRepo: The Hugging Face model repository identifier. + /// - fileName: The path to the `.litertlm` file within the repository. /// - revision: Git revision or branch. Defaults to `main`. - /// - modalities: Defaults to text-only (`[]`) — the safe choice for an - /// unknown model. Pass `.textImage` or `.all` if the model ships those - /// encoders. - /// - visionBackend: Backend for the vision encoder. Defaults to `.cpu()`. - /// - audioBackend: Backend for the audio encoder. Defaults to `.cpu()`. - /// - visualTokenBudget: Per-image visual-token cap (`nil` = engine default). + /// - backend: Backend for text generation. Defaults to Metal GPU. + /// - visionBackend: Backend for the vision encoder; `nil` disables vision. + /// - audioBackend: Backend for the audio encoder; `nil` disables audio. /// - maxTokens: Context (KV cache) budget. - /// - storageDirectory: Where to keep the downloaded model. - /// - onDownloadProgress: Called during the first-run model download. + /// - hub: Optional Hub client for authentication and cache configuration. + /// - downloadProgress: Optional progress object for the model download. public init( huggingFaceRepo: String, fileName: String, revision: String = "main", - modalities: Modality = [], - visionBackend: Backend = .cpu(), - audioBackend: Backend = .cpu(), - visualTokenBudget: Int32? = nil, + backend: Backend = .gpu, + visionBackend: Backend? = nil, + audioBackend: Backend? = nil, maxTokens: Int = 2048, - storageDirectory: URL? = nil, - onDownloadProgress: (@Sendable (ModelDownloader.Progress) -> Void)? = nil + hub: HubClient? = nil, + downloadProgress: Progress? = nil ) { self.engine = LazyEngine { - guard - let url = URL( - string: - "https://huggingface.co/\(huggingFaceRepo)/resolve/\(revision)/\(fileName)?download=true" - ) - else { - throw LiteRTChatError.modelFileNotFound(URL(fileURLWithPath: fileName)) + guard let repo = Repo.ID(rawValue: huggingFaceRepo) else { + throw URLError(.badURL) } - let directory = try storageDirectory ?? LiteRTChat.defaultStorageDirectory() - let destination = directory.appendingPathComponent(fileName) - try await ModelDownloader.shared.download( - from: url, - to: destination, - expectedBytes: nil, - onProgress: onDownloadProgress + let destination = try await (hub ?? .default).downloadFile( + at: fileName, + from: repo, + revision: revision, + progress: downloadProgress ) return try await makeEngine( modelPath: destination.path, - modalities: modalities, + backend: backend, visionBackend: visionBackend, audioBackend: audioBackend, - maxTokens: maxTokens, - visualTokenBudget: visualTokenBudget + maxTokens: maxTokens ) } } @@ -366,27 +306,19 @@ import Foundation private func makeEngine( modelPath: String, - modalities: Modality, - visionBackend: Backend, - audioBackend: Backend, - maxTokens: Int, - visualTokenBudget: Int32? + backend: Backend, + visionBackend: Backend?, + audioBackend: Backend?, + maxTokens: Int ) async throws -> Engine { - ExperimentalFlags.optIntoExperimentalAPIs() - if modalities.contains(.vision), let visualTokenBudget { - ExperimentalFlags.visualTokenBudget = visualTokenBudget - } let caches = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask).first let config = try EngineConfig( modelPath: modelPath, - backend: .gpu, - visionBackend: modalities.contains(.vision) ? visionBackend : nil, - audioBackend: modalities.contains(.audio) ? audioBackend : nil, + backend: backend, + visionBackend: visionBackend, + audioBackend: audioBackend, maxNumTokens: maxTokens, - cacheDir: caches?.path, - // The engine default is 1 image per conversation (a 2nd image - // overwrites the 1st); allow several so multi-image chats work. - maxNumImages: modalities.contains(.vision) ? 16 : nil + cacheDir: caches?.path ) let engine = Engine(engineConfig: config) try await engine.initialize() diff --git a/Tests/AnyLanguageModelTests/LiteRTLanguageModelTests.swift b/Tests/AnyLanguageModelTests/LiteRTLanguageModelTests.swift index dd9a92a5..3ffea57c 100644 --- a/Tests/AnyLanguageModelTests/LiteRTLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/LiteRTLanguageModelTests.swift @@ -3,7 +3,40 @@ import Testing @testable import AnyLanguageModel -#if LiteRT +#if LiteRT && (os(iOS) || os(macOS)) && !targetEnvironment(macCatalyst) + @Suite("LiteRTLanguageModel configuration") + struct LiteRTLanguageModelConfigurationTests { + @Test func missingLocalModelThrows() async { + let url = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString) + .appendingPathExtension("litertlm") + let session = LanguageModelSession(model: LiteRTLanguageModel(modelFileURL: url)) + + do { + _ = try await session.respond(to: "Hello") + Issue.record("Expected a missing model error") + } catch let error as CocoaError { + #expect(error.code == .fileReadNoSuchFile) + } catch { + Issue.record("Unexpected error: \(error)") + } + } + + @Test func invalidRepositoryThrows() async { + let model = LiteRTLanguageModel(huggingFaceRepo: "", fileName: "model.litertlm") + let session = LanguageModelSession(model: model) + + do { + _ = try await session.respond(to: "Hello") + Issue.record("Expected an invalid repository error") + } catch let error as URLError { + #expect(error.code == .badURL) + } catch { + Issue.record("Unexpected error: \(error)") + } + } + } + /// Path to a local `.litertlm` file to test against /// (for example, gemma-4-E2B-it.litertlm). /// These tests load multi-GB weights, so they only run when explicitly @@ -21,8 +54,7 @@ import Testing struct LiteRTLanguageModelTests { private var model: LiteRTLanguageModel { LiteRTLanguageModel( - modelFileURL: URL(fileURLWithPath: liteRTTestModelPath!), - modalities: [] + modelFileURL: URL(fileURLWithPath: liteRTTestModelPath!) ) } From fbd345deee146dab493882bff9bde697f2ea5be1 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Fri, 11 Sep 2026 04:27:46 -0700 Subject: [PATCH 3/6] Use semantic line breaks in LiteRT documentation --- README.md | 29 +++++--- .../Models/LiteRTLanguageModel.swift | 69 +++++++++++-------- .../LiteRTLanguageModelTests.swift | 5 +- 3 files changed, 62 insertions(+), 41 deletions(-) diff --git a/README.md b/README.md index a1527751..b74539a6 100644 --- a/README.md +++ b/README.md @@ -73,7 +73,8 @@ session.toolExecutionDelegate = ToolExecutionObserver() - [x] [Core ML](https://developer.apple.com/documentation/coreml) models - [x] [MLX](https://github.com/ml-explore/mlx-swift) models - [x] [llama.cpp](https://github.com/ggml-org/llama.cpp) (GGUF models) -- [x] [LiteRT-LM](https://github.com/google-ai-edge/litert-lm) (`.litertlm` models, using the official Swift package) +- [x] [LiteRT-LM](https://github.com/google-ai-edge/litert-lm) + (`.litertlm` models, using the official Swift package) - [x] Ollama [HTTP API](https://github.com/ollama/ollama/blob/main/docs/api.md) - [x] Anthropic [Messages API](https://docs.claude.com/en/api/messages) - [x] Google [Gemini API](https://ai.google.dev/api/generate-content) @@ -119,7 +120,8 @@ This results in smaller binary sizes and faster build times. - `Llama`: Enables llama.cpp support (requires `mattt/llama.swift`) - `LiteRT`: Enables LiteRT-LM support for Gemma 4 and other `.litertlm` models - (requires the official `google-ai-edge/LiteRT-LM` Swift package; iOS and macOS only) + (requires the official `google-ai-edge/LiteRT-LM` Swift package; + iOS and macOS only) By default, no traits are enabled. To enable specific traits, specify them in your package's dependencies: @@ -616,9 +618,11 @@ let response = try await session.respond( ### LiteRT-LM -Runs `.litertlm` models (for example, Gemma 4) fully on-device via Google's -[LiteRT-LM](https://github.com/google-ai-edge/litert-lm) runtime with Metal GPU -acceleration (requires `LiteRT` trait; iOS and macOS only): +Runs `.litertlm` models (for example, Gemma 4) fully on-device +via Google's [LiteRT-LM](https://github.com/google-ai-edge/litert-lm) runtime +with Metal GPU acceleration +(requires `LiteRT` trait; +iOS and macOS only): ```swift let model = LiteRTLanguageModel(modelFileURL: modelURL) @@ -627,8 +631,9 @@ let session = LanguageModelSession(model: model) let response = try await session.respond(to: "What is the capital of France?") ``` -You can also load a `.litertlm` file from Hugging Face. The file is downloaded -on first use using the Hub client's cache and authentication: +You can also load a `.litertlm` file from Hugging Face. +The file is downloaded on first use +using the Hub client's cache and authentication: ```swift // Any Hugging Face repo @@ -637,10 +642,12 @@ let model = LiteRTLanguageModel( fileName: "gemma-4-E4B-it.litertlm") ``` -Image inputs are supported for models that ship a vision tower (pass -`visionBackend: .cpu()`). Structured generation is prompt-driven — the JSON -schema is included in the prompt and the response is parsed. Tool calling is -supported for `respond` (not yet for streaming). +Image inputs are supported for models that ship a vision tower +(pass `visionBackend: .cpu()`). +Structured generation is prompt-driven — +the JSON schema is included in the prompt and the response is parsed. +Tool calling is supported for `respond` +(not yet for streaming). ### Ollama diff --git a/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift b/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift index 0ccd1c7b..537a72b9 100644 --- a/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift @@ -5,12 +5,12 @@ import Foundation import class HuggingFace.HubClient import enum HuggingFace.Repo - /// A language model that runs `.litertlm` models fully on-device via Google's - /// [LiteRT-LM](https://github.com/google-ai-edge/litert-lm) runtime. + /// A language model that runs `.litertlm` models fully on-device + /// via Google's [LiteRT-LM](https://github.com/google-ai-edge/litert-lm) runtime. /// - /// Use this model to run Gemma 4 (and other LiteRT-LM models) on iOS and macOS - /// with Metal GPU acceleration, including image understanding for models that - /// ship a vision tower. + /// Use this model to run Gemma 4 (and other LiteRT-LM models) + /// on iOS and macOS with Metal GPU acceleration, + /// including image understanding for models that ship a vision tower. /// /// ```swift /// let model = LiteRTLanguageModel(modelFileURL: modelURL) @@ -19,16 +19,19 @@ import Foundation /// ``` /// /// Hugging Face models are downloaded on first use and stored in the Hub cache. - /// Loading is lazy: the engine is brought up on the first - /// request (or when ``prewarm(for:promptPrefix:)`` is called). + /// Loading is lazy: + /// the engine is brought up on the first request + /// (or when ``prewarm(for:promptPrefix:)`` is called). /// - /// Structured generation is prompt-driven (the JSON schema is included in the - /// prompt and the response is parsed), and tool calling is supported for + /// Structured generation is prompt-driven + /// (the JSON schema is included in the prompt and the response is parsed), + /// and tool calling is supported for /// ``respond(within:to:generating:includeSchemaInPrompt:options:)`` /// (not yet for streaming). public struct LiteRTLanguageModel: LanguageModel { /// The reason the model is unavailable. - /// This model is always available; loading errors surface when responding. + /// This model is always available; + /// loading errors surface when responding. public typealias UnavailableReason = Never private let engine: LazyEngine @@ -37,10 +40,13 @@ import Foundation /// /// - Parameters: /// - modelFileURL: File URL of an on-disk `.litertlm` model. - /// - backend: Backend for text generation. Defaults to Metal GPU. - /// - visionBackend: Backend for the vision encoder. Pass `.cpu()` for - /// a model with image support; `nil` disables vision. - /// - audioBackend: Backend for the audio encoder; `nil` disables audio. + /// - backend: Backend for text generation. + /// Defaults to Metal GPU. + /// - visionBackend: Backend for the vision encoder. + /// Pass `.cpu()` for a model with image support; + /// `nil` disables vision. + /// - audioBackend: Backend for the audio encoder; + /// `nil` disables audio. /// - maxTokens: Context (KV cache) budget. public init( modelFileURL: URL, @@ -71,10 +77,14 @@ import Foundation /// - Parameters: /// - huggingFaceRepo: The Hugging Face model repository identifier. /// - fileName: The path to the `.litertlm` file within the repository. - /// - revision: Git revision or branch. Defaults to `main`. - /// - backend: Backend for text generation. Defaults to Metal GPU. - /// - visionBackend: Backend for the vision encoder; `nil` disables vision. - /// - audioBackend: Backend for the audio encoder; `nil` disables audio. + /// - revision: Git revision or branch. + /// Defaults to `main`. + /// - backend: Backend for text generation. + /// Defaults to Metal GPU. + /// - visionBackend: Backend for the vision encoder; + /// `nil` disables vision. + /// - audioBackend: Backend for the audio encoder; + /// `nil` disables audio. /// - maxTokens: Context (KV cache) budget. /// - hub: Optional Hub client for authentication and cache configuration. /// - downloadProgress: Optional progress object for the model download. @@ -286,7 +296,8 @@ import Foundation // MARK: - Engine Bring-Up /// Brings up the engine on first use and shares it across requests. - /// Engine bring-up loads multi-GB weights, so it must happen exactly once. + /// Engine bring-up loads multi-GB weights, + /// so it must happen exactly once. private actor LazyEngine { private var task: Task? private let bringUp: @Sendable () async throws -> Engine @@ -332,9 +343,9 @@ import Foundation var history: [Message] var prompt: Message - /// Extends the plan after a tool round-trip: the trigger prompt and the - /// model's tool-call text become history, and the tool result becomes the - /// new trigger. + /// Extends the plan after a tool round-trip: + /// the trigger prompt and the model's tool-call text become history, + /// and the tool result becomes the new trigger. func continuing(afterModelText text: String, toolOutput: Transcript.ToolOutput) -> GenerationPlan { var history = self.history history.append(prompt) @@ -348,9 +359,10 @@ import Foundation } } - /// Splits the session transcript into a system message, prior turns, and the - /// message to generate from. The generation trigger is the last `.prompt` or - /// (in a tool round-trip) the last `.toolOutput` entry. + /// Splits the session transcript into a system message, + /// prior turns, and the message to generate from. + /// The generation trigger is the last `.prompt` + /// or (in a tool round-trip) the last `.toolOutput` entry. private func makePlan( from transcript: Transcript, fallbackPrompt: String, @@ -411,8 +423,8 @@ import Foundation ) } - /// Maps transcript segments to LiteRT content: text, structured content as - /// JSON text, and images. + /// Maps transcript segments to LiteRT content: + /// text, structured content as JSON text, and images. private func messageContents(of segments: [Transcript.Segment]) -> [Content] { var contents: [Content] = [] for segment in segments { @@ -528,7 +540,8 @@ import Foundation return lines.joined(separator: "\n") } - /// Parses a tool call from model output, if present and naming a known tool. + /// Parses a tool call from model output, + /// if present and naming a known tool. private func parseToolCall( from text: String, tools: [any Tool] diff --git a/Tests/AnyLanguageModelTests/LiteRTLanguageModelTests.swift b/Tests/AnyLanguageModelTests/LiteRTLanguageModelTests.swift index 3ffea57c..6c45cdda 100644 --- a/Tests/AnyLanguageModelTests/LiteRTLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/LiteRTLanguageModelTests.swift @@ -39,8 +39,9 @@ import Testing /// Path to a local `.litertlm` file to test against /// (for example, gemma-4-E2B-it.litertlm). - /// These tests load multi-GB weights, so they only run when explicitly - /// requested via the `LITERT_TEST_MODEL` environment variable. + /// These tests load multi-GB weights, + /// so they only run when explicitly requested + /// via the `LITERT_TEST_MODEL` environment variable. private let liteRTTestModelPath = ProcessInfo.processInfo.environment["LITERT_TEST_MODEL"] private let shouldRunLiteRTTests = liteRTTestModelPath != nil From 375fc891b1aefee5f42889a63f2e38371b70a1b2 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Fri, 11 Sep 2026 05:31:50 -0700 Subject: [PATCH 4/6] Address LiteRT generation and loading review feedback --- .github/workflows/ci.yml | 9 +- Package.resolved | 200 +++++++++- .../Models/LiteRTLanguageModel.swift | 120 +++--- .../Models/LiteRTRuntime.swift | 131 ++++++ .../LiteRTLanguageModelBehaviorTests.swift | 374 ++++++++++++++++++ 5 files changed, 785 insertions(+), 49 deletions(-) create mode 100644 Sources/AnyLanguageModel/Models/LiteRTRuntime.swift create mode 100644 Tests/AnyLanguageModelTests/LiteRTLanguageModelBehaviorTests.swift diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 82cfaa96..25f77fbd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -10,6 +10,11 @@ permissions: contents: read pull-requests: write +# LiteRT-LM's Swift package uses release XCFrameworks; +# its Git LFS assets are not needed by SwiftPM. +env: + GIT_LFS_SKIP_SMUDGE: "1" + jobs: test-macos: name: Swift ${{ matrix.swift }} on macOS ${{ matrix.macos }} with Xcode ${{ matrix.xcode }}${{ matrix.traits != '' && format(' and --traits {0}', matrix.traits) || '' }} @@ -63,10 +68,10 @@ jobs: run: swift format lint --strict --recursive . - name: Build - run: swift build --build-tests --traits MLX,Llama,CoreML${{ matrix.traits != '' && format(',{0}', matrix.traits) || '' }} + run: swift build --build-tests --traits MLX,Llama,CoreML,LiteRT${{ matrix.traits != '' && format(',{0}', matrix.traits) || '' }} - name: Test - run: swift test --skip-build --traits MLX,Llama,CoreML${{ matrix.traits != '' && format(',{0}', matrix.traits) || '' }} + run: swift test --skip-build --traits MLX,Llama,CoreML,LiteRT${{ matrix.traits != '' && format(',{0}', matrix.traits) || '' }} test-linux: name: Swift ${{ matrix.swift }} on Linux${{ matrix.traits != '' && format(' and --traits {0}', matrix.traits) || '' }} diff --git a/Package.resolved b/Package.resolved index e1984540..5d72bfa5 100644 --- a/Package.resolved +++ b/Package.resolved @@ -1,6 +1,15 @@ { - "originHash" : "802c7ad583bfa4b4de0e754641d0e364438151ba6bdd8bcbcd9050572a0cfc4f", + "originHash" : "872262beb92df73ada611fd9ca3cc641833417deb1e27d47f5df8176d080243d", "pins" : [ + { + "identity" : "async-http-client", + "kind" : "remoteSourceControl", + "location" : "https://github.com/swift-server/async-http-client.git", + "state" : { + "revision" : "f95c908967e98c68c5ce3fd61a7974e7e869e303", + "version" : "1.36.1" + } + }, { "identity" : "eventsource", "kind" : "remoteSourceControl", @@ -37,6 +46,24 @@ "version" : "2.10549.0" } }, + { + "identity" : "mlx-swift", + "kind" : "remoteSourceControl", + "location" : "https://github.com/ml-explore/mlx-swift", + "state" : { + "revision" : "0bb916c67f4b9e5c682cbe02a42c701c93ab5021", + "version" : "0.31.6" + } + }, + { + "identity" : "mlx-swift-lm", + "kind" : "remoteSourceControl", + "location" : "https://github.com/ml-explore/mlx-swift-lm", + "state" : { + "revision" : "bd4b7434e6bdb588c7ef55706ff8904cb7fd4c57", + "version" : "3.31.4" + } + }, { "identity" : "partialjsondecoder", "kind" : "remoteSourceControl", @@ -46,6 +73,24 @@ "version" : "1.0.0" } }, + { + "identity" : "swift-algorithms", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-algorithms.git", + "state" : { + "revision" : "87e50f483c54e6efd60e885f7f5aa946cee68023", + "version" : "1.2.1" + } + }, + { + "identity" : "swift-argument-parser", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-argument-parser", + "state" : { + "revision" : "6a52f3251125d74daf04fcbd5e6f08a75d074382", + "version" : "1.8.2" + } + }, { "identity" : "swift-asn1", "kind" : "remoteSourceControl", @@ -55,6 +100,15 @@ "version" : "1.7.2" } }, + { + "identity" : "swift-async-algorithms", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-async-algorithms.git", + "state" : { + "revision" : "3da39bbc4e687d4192af7c9cf4eab805745a0b9c", + "version" : "1.1.5" + } + }, { "identity" : "swift-atomics", "kind" : "remoteSourceControl", @@ -64,6 +118,15 @@ "version" : "1.3.1" } }, + { + "identity" : "swift-certificates", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-certificates.git", + "state" : { + "revision" : "c8aece90ea05f9866bd392a5bf13b5cae56c0e03", + "version" : "1.20.0" + } + }, { "identity" : "swift-collections", "kind" : "remoteSourceControl", @@ -73,6 +136,15 @@ "version" : "1.6.0" } }, + { + "identity" : "swift-configuration", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-configuration.git", + "state" : { + "revision" : "be76c4ad929eb6c4bcaf3351799f2adf9e6848a9", + "version" : "1.2.0" + } + }, { "identity" : "swift-crypto", "kind" : "remoteSourceControl", @@ -82,6 +154,33 @@ "version" : "4.5.2" } }, + { + "identity" : "swift-distributed-tracing", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-distributed-tracing.git", + "state" : { + "revision" : "dc4030184203ffafbb2ec614352487235d747fe0", + "version" : "1.4.1" + } + }, + { + "identity" : "swift-http-structured-headers", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-http-structured-headers.git", + "state" : { + "revision" : "933538faa42c432d385f02e07df0ace7c5ecfc47", + "version" : "1.7.0" + } + }, + { + "identity" : "swift-http-types", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-http-types.git", + "state" : { + "revision" : "bff4b6903cdc99dda49649dd52f46c11cfd3ed50", + "version" : "1.8.0" + } + }, { "identity" : "swift-huggingface", "kind" : "remoteSourceControl", @@ -91,6 +190,24 @@ "version" : "0.10.1" } }, + { + "identity" : "swift-jinja", + "kind" : "remoteSourceControl", + "location" : "https://github.com/huggingface/swift-jinja.git", + "state" : { + "revision" : "12ea4955e380dae6b0cd98299effe48c7fe584fc", + "version" : "2.5.0" + } + }, + { + "identity" : "swift-log", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-log.git", + "state" : { + "revision" : "9c6fb14227f55d8f711ce3847dc2f419fb0ecacb", + "version" : "1.15.1" + } + }, { "identity" : "swift-nio", "kind" : "remoteSourceControl", @@ -100,6 +217,69 @@ "version" : "2.101.2" } }, + { + "identity" : "swift-nio-extras", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-nio-extras.git", + "state" : { + "revision" : "41449336c8ecfadac6b4b5be75f9c3c306e61ced", + "version" : "1.35.1" + } + }, + { + "identity" : "swift-nio-http2", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-nio-http2.git", + "state" : { + "revision" : "0f3e54e29c944c2e835ad52159da7d9e1c94ac69", + "version" : "1.46.0" + } + }, + { + "identity" : "swift-nio-ssl", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-nio-ssl.git", + "state" : { + "revision" : "03827c1a9fdb2b6b00a4e93ede8861520263af8c", + "version" : "2.37.4" + } + }, + { + "identity" : "swift-nio-transport-services", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-nio-transport-services.git", + "state" : { + "revision" : "67787bb645a5e67d2edcdfbe48a216cc549222d5", + "version" : "1.28.0" + } + }, + { + "identity" : "swift-numerics", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-numerics", + "state" : { + "revision" : "0c0290ff6b24942dadb83a929ffaaa1481df04a2", + "version" : "1.1.1" + } + }, + { + "identity" : "swift-service-context", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-service-context.git", + "state" : { + "revision" : "d0997351b0c7779017f88e7a93bc30a1878d7f29", + "version" : "1.3.0" + } + }, + { + "identity" : "swift-service-lifecycle", + "kind" : "remoteSourceControl", + "location" : "https://github.com/swift-server/swift-service-lifecycle", + "state" : { + "revision" : "7f9326b0326ff86e3646295ea6e891f68c471c5e", + "version" : "2.12.0" + } + }, { "identity" : "swift-syntax", "kind" : "remoteSourceControl", @@ -117,6 +297,24 @@ "revision" : "7502b711c92a17741fa625d722b0ccbd595d8ed1", "version" : "1.7.2" } + }, + { + "identity" : "swift-transformers", + "kind" : "remoteSourceControl", + "location" : "https://github.com/huggingface/swift-transformers", + "state" : { + "revision" : "c21fdcde390313a6d98d8e33a346f2c3486c3ab0", + "version" : "1.3.4" + } + }, + { + "identity" : "yyjson", + "kind" : "remoteSourceControl", + "location" : "https://github.com/ibireme/yyjson.git", + "state" : { + "revision" : "8b4a38dc994a110abaec8a400615567bd996105f", + "version" : "0.12.0" + } } ], "version" : 3 diff --git a/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift b/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift index 537a72b9..5e321e79 100644 --- a/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift @@ -34,7 +34,8 @@ import Foundation /// loading errors surface when responding. public typealias UnavailableReason = Never - private let engine: LazyEngine + private let engine: LiteRTModelLoader + private let imageSession: URLSession /// Creates a model from a local `.litertlm` file. /// @@ -55,7 +56,8 @@ import Foundation audioBackend: Backend? = nil, maxTokens: Int = 2048 ) { - self.engine = LazyEngine { + self.imageSession = .shared + self.engine = LiteRTModelLoader { guard modelFileURL.isFileURL, FileManager.default.fileExists(atPath: modelFileURL.path) else { @@ -99,7 +101,8 @@ import Foundation hub: HubClient? = nil, downloadProgress: Progress? = nil ) { - self.engine = LazyEngine { + self.imageSession = .shared + self.engine = LiteRTModelLoader { guard let repo = Repo.ID(rawValue: huggingFaceRepo) else { throw URLError(.badURL) } @@ -119,6 +122,14 @@ import Foundation } } + init( + load: @escaping @Sendable () async throws -> any LiteRTRuntime, + imageSession: URLSession = .shared + ) { + self.engine = LiteRTModelLoader(load) + self.imageSession = imageSession + } + public func prewarm( for session: LanguageModelSession, promptPrefix: Prompt? @@ -144,11 +155,12 @@ import Foundation } let tools = session.tools - var plan = makePlan( + var plan = try await makePlan( from: session.transcript, fallbackPrompt: prompt.description, schemaJSON: includeSchemaInPrompt ? schemaJSON : nil, - tools: tools + tools: tools, + imageSession: imageSession ) let sampler = makeSampler(for: options, structured: schemaJSON != nil || !tools.isEmpty) @@ -157,8 +169,8 @@ import Foundation var toolRounds = 0 while true { - let conversation = try await engine.createConversation( - with: ConversationConfig( + let conversation = try await engine.makeConversation( + config: ConversationConfig( systemMessage: plan.systemMessage, initialMessages: plan.history, samplerConfig: sampler @@ -166,14 +178,22 @@ import Foundation ) text = "" - for try await chunk in conversation.sendMessageStream(plan.prompt) { - text += chunk.toString + try await generateLiteRTResponse( + conversation: conversation, + prompt: plan.prompt, + maximumResponseTokens: options.maximumResponseTokens + ) { chunk in + text += chunk } guard !tools.isEmpty, - toolRounds < maxToolRounds, let parsed = parseToolCall(from: text, tools: tools) else { break } + guard toolRounds < maxToolRounds else { + throw LanguageModelSession.GenerationError.decodingFailure( + .init(debugDescription: "Exceeded maximum LiteRT tool iterations (\(maxToolRounds)).") + ) + } toolRounds += 1 let resolution = try await resolveToolCall( @@ -183,6 +203,14 @@ import Foundation ) switch resolution { case .stop(let call): + guard type == String.self else { + throw LanguageModelSession.GenerationError.decodingFailure( + .init( + debugDescription: + "Tool execution stopped before LiteRT generated a structured response." + ) + ) + } entries.append(.toolCalls(Transcript.ToolCalls([call]))) return LanguageModelSession.Response( content: "" as! Content, @@ -235,14 +263,15 @@ import Foundation schemaJSON = try encodeSchema(type.generationSchema) } - let plan = makePlan( + let plan = try await makePlan( from: session.transcript, fallbackPrompt: prompt.description, schemaJSON: includeSchemaInPrompt ? schemaJSON : nil, - tools: [] + tools: [], + imageSession: imageSession ) - let conversation = try await engine.createConversation( - with: ConversationConfig( + let conversation = try await engine.makeConversation( + config: ConversationConfig( systemMessage: plan.systemMessage, initialMessages: plan.history, samplerConfig: makeSampler(for: options, structured: schemaJSON != nil) @@ -250,9 +279,12 @@ import Foundation ) var text = "" - for try await chunk in conversation.sendMessageStream(plan.prompt) { - let delta = chunk.toString - guard !delta.isEmpty else { continue } + try await generateLiteRTResponse( + conversation: conversation, + prompt: plan.prompt, + maximumResponseTokens: options.maximumResponseTokens + ) { delta in + guard !delta.isEmpty else { return } text += delta if type == String.self { @@ -295,26 +327,6 @@ import Foundation // MARK: - Engine Bring-Up - /// Brings up the engine on first use and shares it across requests. - /// Engine bring-up loads multi-GB weights, - /// so it must happen exactly once. - private actor LazyEngine { - private var task: Task? - private let bringUp: @Sendable () async throws -> Engine - - init(_ bringUp: @escaping @Sendable () async throws -> Engine) { - self.bringUp = bringUp - } - - func ready() async throws -> Engine { - if task == nil { - let bringUp = self.bringUp - task = Task { try await bringUp() } - } - return try await task!.value - } - } - private func makeEngine( modelPath: String, backend: Backend, @@ -367,8 +379,9 @@ import Foundation from transcript: Transcript, fallbackPrompt: String, schemaJSON: String?, - tools: [any Tool] - ) -> GenerationPlan { + tools: [any Tool], + imageSession: URLSession + ) async throws -> GenerationPlan { let entries = Array(transcript) let triggerIndex = entries.lastIndex { entry in switch entry { @@ -378,8 +391,9 @@ import Foundation } var systemText: [String] = [] - if !tools.isEmpty { - systemText.append(toolInstructions(tools)) + let describedTools = tools.filter(\.includesSchemaInInstructions) + if !describedTools.isEmpty { + systemText.append(toolInstructions(describedTools)) } var history: [Message] = [] var trigger: Message? @@ -390,7 +404,7 @@ import Foundation case .instructions(let instructions): systemText.append(textContent(of: instructions.segments)) case .prompt(let prompt): - var contents = messageContents(of: prompt.segments) + var contents = try await messageContents(of: prompt.segments, session: imageSession) if isTrigger, let schemaJSON, !schemaJSON.isEmpty { contents.append( .text( @@ -425,7 +439,10 @@ import Foundation /// Maps transcript segments to LiteRT content: /// text, structured content as JSON text, and images. - private func messageContents(of segments: [Transcript.Segment]) -> [Content] { + private func messageContents( + of segments: [Transcript.Segment], + session: URLSession + ) async throws -> [Content] { var contents: [Content] = [] for segment in segments { switch segment { @@ -442,7 +459,13 @@ import Foundation case .url(let url): if url.isFileURL { contents.append(.imageFile(url.path)) - } else if let data = try? Data(contentsOf: url) { + } else { + let (data, response) = try await session.data(from: url) + if let response = response as? HTTPURLResponse, + !(200 ... 299).contains(response.statusCode) + { + throw URLError(.badServerResponse) + } contents.append(.imageData(data)) } } @@ -451,10 +474,15 @@ import Foundation return contents.isEmpty ? [.text("")] : contents } - /// Concatenates the text of a segment list (non-text segments are ignored). + /// Concatenates text and structured JSON from a segment list. + /// Image segments are ignored. private func textContent(of segments: [Transcript.Segment]) -> String { segments.compactMap { segment in - if case .text(let text) = segment { return text.content } else { return nil } + switch segment { + case .text(let text): return text.content + case .structure(let structure): return structure.content.jsonString + case .image: return nil + } }.joined(separator: " ") } diff --git a/Sources/AnyLanguageModel/Models/LiteRTRuntime.swift b/Sources/AnyLanguageModel/Models/LiteRTRuntime.swift new file mode 100644 index 00000000..401f0dea --- /dev/null +++ b/Sources/AnyLanguageModel/Models/LiteRTRuntime.swift @@ -0,0 +1,131 @@ +import Foundation + +#if LiteRT && (os(iOS) || os(macOS)) && !targetEnvironment(macCatalyst) + @preconcurrency import LiteRTLM + + protocol LiteRTRuntime: Sendable { + func makeConversation(config: ConversationConfig) async throws -> any LiteRTConversation + } + + protocol LiteRTConversation: Sendable { + func sendMessageStream(_ message: Message, maxOutputTokens: Int?) -> AsyncThrowingStream + func cancel() throws + } + + extension Engine: LiteRTRuntime { + func makeConversation(config: ConversationConfig) throws -> any LiteRTConversation { + NativeLiteRTConversation(conversation: try createConversation(with: config)) + } + } + + private struct NativeLiteRTConversation: LiteRTConversation { + let conversation: Conversation + + func sendMessageStream(_ message: Message, maxOutputTokens: Int?) -> AsyncThrowingStream { + conversation.sendMessageStream(message, maxOutputTokens: maxOutputTokens) + } + + func cancel() throws { + try conversation.cancel() + } + } + + /// Shares successful loading and in-flight work across requests. + /// A failed attempt is cleared only by its own waiters, + /// so an older failure cannot discard a newer retry. + actor LiteRTModelLoader { + private var attempt: (id: UUID, task: Task)? + private let load: @Sendable () async throws -> any LiteRTRuntime + + init(_ load: @escaping @Sendable () async throws -> any LiteRTRuntime) { + self.load = load + } + + func ready() async throws -> any LiteRTRuntime { + try Task.checkCancellation() + if attempt == nil { + let load = self.load + attempt = (UUID(), Task { try await load() }) + } + let current = attempt! + let runtime: any LiteRTRuntime + do { + runtime = try await current.task.value + } catch { + if attempt?.id == current.id { + attempt = nil + } + throw error + } + try Task.checkCancellation() + return runtime + } + } + + /// Serializes native inference startup with cancellation. + /// Cancelling before startup prevents inference from starting; + /// cancelling after startup stops the active conversation. + private final class LiteRTGeneration: Sendable { + private struct State: Sendable { + var isCancelled = false + var conversation: (any LiteRTConversation)? + } + + private let state = Locked(State()) + + func start( + conversation: any LiteRTConversation, + prompt: Message, + maximumResponseTokens: Int? + ) throws -> AsyncThrowingStream { + try state.withLock { state in + guard !state.isCancelled else { throw CancellationError() } + let stream = conversation.sendMessageStream(prompt, maxOutputTokens: maximumResponseTokens) + state.conversation = conversation + return stream + } + } + + func cancel() { + let conversation = state.withLock { state in + state.isCancelled = true + let conversation = state.conversation + state.conversation = nil + return conversation + } + try? conversation?.cancel() + } + + func finish() { + state.withLock { $0.conversation = nil } + } + } + + func generateLiteRTResponse( + conversation: any LiteRTConversation, + prompt: Message, + maximumResponseTokens: Int?, + onChunk: (String) throws -> Void + ) async throws { + let generation = LiteRTGeneration() + try await withTaskCancellationHandler { + defer { + if Task.isCancelled { generation.cancel() } + generation.finish() + } + try Task.checkCancellation() + let stream = try generation.start( + conversation: conversation, + prompt: prompt, + maximumResponseTokens: maximumResponseTokens + ) + for try await chunk in stream { + try Task.checkCancellation() + try onChunk(chunk.toString) + } + try Task.checkCancellation() + } onCancel: { + generation.cancel() + } + } +#endif diff --git a/Tests/AnyLanguageModelTests/LiteRTLanguageModelBehaviorTests.swift b/Tests/AnyLanguageModelTests/LiteRTLanguageModelBehaviorTests.swift new file mode 100644 index 00000000..d794108d --- /dev/null +++ b/Tests/AnyLanguageModelTests/LiteRTLanguageModelBehaviorTests.swift @@ -0,0 +1,374 @@ +import Foundation +import Testing + +@testable import AnyLanguageModel + +#if LiteRT && (os(iOS) || os(macOS)) && !targetEnvironment(macCatalyst) + @preconcurrency import LiteRTLM + + @Suite("LiteRTLanguageModel behavior", .timeLimit(.minutes(1))) + struct LiteRTLanguageModelBehaviorTests { + @Test(arguments: [false, true]) + func forwardsResponseTokenLimit(streaming: Bool) async throws { + let conversation = StubLiteRTConversation(reply: "Hello") + let runtime = StubLiteRTRuntime([conversation]) + let session = LanguageModelSession(model: LiteRTLanguageModel(load: { runtime })) + let options = GenerationOptions(maximumResponseTokens: 7) + if streaming { + for try await _ in session.streamResponse(to: "Hello", options: options) {} + } else { + _ = try await session.respond(to: "Hello", options: options) + } + #expect(conversation.requests.withLock { $0.map(\.maxOutputTokens) } == [7]) + #expect(conversation.cancelCount.withLock { $0 } == 0) + } + + @Test(arguments: [false, true]) + func cancelsNativeInference(streaming: Bool) async throws { + let conversation = StubLiteRTConversation(reply: "Hello", keepRunning: true) + let runtime = StubLiteRTRuntime([conversation]) + let session = LanguageModelSession(model: LiteRTLanguageModel(load: { runtime })) + let task = Task { + if streaming { + for try await _ in session.streamResponse(to: "Hello") {} + } else { + _ = try await session.respond(to: "Hello") + } + } + await conversation.started.wait() + task.cancel() + if streaming { + do { try await task.value } catch { #expect(error is CancellationError) } + } else { + await #expect(throws: CancellationError.self) { try await task.value } + } + await conversation.cancelled.wait() + #expect(conversation.cancelCount.withLock { $0 } == 1) + } + + @Test func abandoningStreamCancelsNativeInference() async throws { + let conversation = StubLiteRTConversation(reply: "Hello", keepRunning: true) + let runtime = StubLiteRTRuntime([conversation]) + let session = LanguageModelSession(model: LiteRTLanguageModel(load: { runtime })) + for try await _ in session.streamResponse(to: "Hello") { break } + await conversation.cancelled.wait() + #expect(conversation.cancelCount.withLock { $0 } == 1) + } + + @Test func cancellationDuringLoadingDoesNotStartInference() async throws { + let started = LiteRTTestSignal() + let release = LiteRTTestSignal() + let runtime = StubLiteRTRuntime([StubLiteRTConversation(reply: "Hello")]) + let model = LiteRTLanguageModel(load: { + await started.signal() + await release.wait() + return runtime + }) + let session = LanguageModelSession(model: model) + let task = Task { try await session.respond(to: "Hello") } + await started.wait() + task.cancel() + await release.signal() + await #expect(throws: CancellationError.self) { try await task.value } + #expect(runtime.configurations.withLock { $0.isEmpty }) + _ = try await session.respond(to: "Try again") + #expect(runtime.configurations.withLock { $0.count } == 1) + } + + @Test func retriesFailedLoadingAndSharesSuccessfulRuntime() async throws { + let loads = Locked(0) + let runtime = StubLiteRTRuntime([ + StubLiteRTConversation(reply: "Hello"), + StubLiteRTConversation(reply: "Again"), + ]) + let model = LiteRTLanguageModel(load: { + let count = loads.withLock { + $0 += 1; return $0 + } + if count == 1 { throw URLError(.timedOut) } + return runtime + }) + let session = LanguageModelSession(model: model) + await #expect(throws: URLError.self) { try await session.respond(to: "Hello") } + #expect(try await session.respond(to: "Retry").content == "Hello") + #expect(try await session.respond(to: "Again").content == "Again") + #expect(loads.withLock { $0 } == 2) + } + + @Test func concurrentWaitersCannotDiscardANewerRetry() async throws { + let loads = Locked(0) + let started = LiteRTTestSignal() + let release = LiteRTTestSignal() + let runtime = StubLiteRTRuntime([]) + let loader = LiteRTModelLoader { + let count = loads.withLock { + $0 += 1; return $0 + } + if count == 1 { + await started.signal() + await release.wait() + throw URLError(.timedOut) + } + return runtime + } + try await withThrowingTaskGroup(of: Void.self) { group in + for _ in 0 ..< 32 { + group.addTask { + _ = try? await loader.ready() + _ = try await loader.ready() + } + } + await started.wait() + await release.signal() + try await group.waitForAll() + } + #expect(loads.withLock { $0 } == 2) + } + + @Test(arguments: [#"{"answer":42}"#, "[1,2]", "42", "true", #""hello""#]) + func preservesStructuredToolOutputAndTranscript(json: String) async throws { + let output = try GeneratedContent(json: json) + let tool = LiteRTTestTool(output: output) + let first = StubLiteRTConversation(reply: toolCall) + let second = StubLiteRTConversation(reply: "Done") + let third = StubLiteRTConversation(reply: "Remembered") + let runtime = StubLiteRTRuntime([first, second, third]) + let session = LanguageModelSession(model: LiteRTLanguageModel(load: { runtime }), tools: [tool]) + _ = try await session.respond(to: "Use the tool", options: .init(maximumResponseTokens: 9)) + let nextPrompt = try #require(second.requests.withLock { $0.first?.text }) + #expect(nextPrompt.contains(output.jsonString)) + #expect(second.requests.withLock { $0.first?.maxOutputTokens } == 9) + _ = try await session.respond(to: "What did the tool return?") + let history = try #require(runtime.configurations.withLock { $0.last?.history }) + #expect(history.contains { $0.contains(output.jsonString) }) + } + + @Test func stoppedStructuredResponseThrowsInsteadOfTrapping() async throws { + let tool = LiteRTTestTool(output: GeneratedContent("unused")) + let runtime = StubLiteRTRuntime([StubLiteRTConversation(reply: toolCall)]) + let session = LanguageModelSession(model: LiteRTLanguageModel(load: { runtime }), tools: [tool]) + session.toolExecutionDelegate = LiteRTStopDelegate() + await #expect { + try await session.respond(to: "Use the tool", generating: LiteRTTestAnswer.self) + } throws: { error in + guard case LanguageModelSession.GenerationError.decodingFailure = error else { return false } + return true + } + #expect(tool.calls.withLock { $0 } == 0) + } + + @Test func stoppedTextResponsePreservesPendingToolCall() async throws { + let tool = LiteRTTestTool(output: GeneratedContent("unused")) + let runtime = StubLiteRTRuntime([StubLiteRTConversation(reply: toolCall)]) + let session = LanguageModelSession(model: LiteRTLanguageModel(load: { runtime }), tools: [tool]) + session.toolExecutionDelegate = LiteRTStopDelegate() + #expect(try await session.respond(to: "Use the tool").content.isEmpty) + #expect( + session.transcript.contains { + if case .toolCalls = $0 { return true }; return false + } + ) + #expect(tool.calls.withLock { $0 } == 0) + } + + @Test func rejectsUnresolvedToolCallAtIterationLimit() async throws { + let tool = LiteRTTestTool(output: GeneratedContent("value")) + let runtime = StubLiteRTRuntime((0 ..< 5).map { _ in StubLiteRTConversation(reply: toolCall) }) + let session = LanguageModelSession(model: LiteRTLanguageModel(load: { runtime }), tools: [tool]) + await #expect { + try await session.respond(to: "Keep calling the tool") + } throws: { error in + guard case LanguageModelSession.GenerationError.decodingFailure = error else { return false } + return true + } + #expect(tool.calls.withLock { $0 } == 4) + } + + @Test func omitsHiddenToolSchemaButStillExecutesItsCalls() async throws { + let tool = LiteRTTestTool(output: GeneratedContent("value"), includesSchemaInInstructions: false) + let runtime = StubLiteRTRuntime([ + StubLiteRTConversation(reply: toolCall), StubLiteRTConversation(reply: "Done"), + ]) + let session = LanguageModelSession( + model: LiteRTLanguageModel(load: { runtime }), + tools: [tool], + instructions: "Help the user" + ) + _ = try await session.respond(to: "Use your built-in tool") + #expect(runtime.configurations.withLock { $0.first?.systemMessage } == "Help the user") + #expect(tool.calls.withLock { $0 } == 1) + } + + @Test(arguments: ["success", "missing", "network"], [false, true]) + func remoteImagesUseAsyncDownloadsAndPropagateErrors(path: String, streaming: Bool) async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [LiteRTImageURLProtocol.self] + let imageSession = URLSession(configuration: configuration) + defer { imageSession.invalidateAndCancel() } + let conversation = StubLiteRTConversation(reply: "An image") + let runtime = StubLiteRTRuntime([conversation]) + let model = LiteRTLanguageModel(load: { runtime }, imageSession: imageSession) + let session = LanguageModelSession(model: model) + let image = Transcript.ImageSegment(url: URL(string: "https://images.invalid/\(path)")!) + do { + if streaming { + for try await _ in session.streamResponse(to: "Describe", images: [image]) {} + } else { + _ = try await session.respond(to: "Describe", images: [image]) + } + #expect(path == "success") + #expect(conversation.requests.withLock { $0.first?.images } == [Data("image bytes".utf8)]) + } catch let error as URLError { + #expect(error.code == (path == "missing" ? .badServerResponse : .notConnectedToInternet)) + #expect(path != "success") + #expect(conversation.requests.withLock { $0.isEmpty }) + } + } + } + + private let toolCall = #"{"tool_call":{"name":"lookup","arguments":{"query":"answer"}}}"# + + @Generable + private struct LiteRTTestAnswer { + var answer: String + } + + private struct LiteRTTestTool: AnyLanguageModel.Tool { + let name = "lookup" + let description = "Looks up an answer" + let output: GeneratedContent + var includesSchemaInInstructions = true + let calls = Locked(0) + + @Generable + struct Arguments { + var query: String + } + + func call(arguments: Arguments) async throws -> GeneratedContent { + calls.withLock { $0 += 1 } + return output + } + } + + private struct LiteRTStopDelegate: ToolExecutionDelegate { + func toolCallDecision( + for toolCall: Transcript.ToolCall, + in session: LanguageModelSession + ) async -> ToolExecutionDecision { .stop } + } + + private actor LiteRTTestSignal { + private var signalled = false + private var waiters: [CheckedContinuation] = [] + + func wait() async { + if signalled { return } + await withCheckedContinuation { waiters.append($0) } + } + + func signal() { + signalled = true + for waiter in waiters { waiter.resume() } + waiters = [] + } + } + + private final class StubLiteRTRuntime: LiteRTRuntime { + struct Configuration: Sendable { + var systemMessage: String? + var history: [String] + } + let configurations = Locked<[Configuration]>([]) + let conversations: [StubLiteRTConversation] + + init(_ conversations: [StubLiteRTConversation]) { + self.conversations = conversations + } + + func makeConversation(config: ConversationConfig) async throws -> any LiteRTConversation { + let index = configurations.withLock { records in + let index = records.count + records.append( + Configuration( + systemMessage: config.systemMessage?.toString, + history: config.initialMessages.map(\.toString) + ) + ) + return index + } + guard conversations.indices.contains(index) else { throw URLError(.resourceUnavailable) } + return conversations[index] + } + } + + private final class StubLiteRTConversation: LiteRTConversation { + struct Request: Sendable { + var text: String + var images: [Data] + var maxOutputTokens: Int? + } + let reply: String + let keepRunning: Bool + let requests = Locked<[Request]>([]) + let cancelCount = Locked(0) + let continuation = Locked.Continuation?>(nil) + let started = LiteRTTestSignal() + let cancelled = LiteRTTestSignal() + + init(reply: String, keepRunning: Bool = false) { + self.reply = reply + self.keepRunning = keepRunning + } + + func sendMessageStream(_ message: Message, maxOutputTokens: Int?) -> AsyncThrowingStream { + requests.withLock { + $0.append( + Request( + text: message.toString, + images: message.contents.compactMap { + if case .imageData(let data) = $0 { return data }; return nil + }, + maxOutputTokens: maxOutputTokens + ) + ) + } + return AsyncThrowingStream { continuation in + self.continuation.withLock { $0 = continuation } + continuation.yield(Message(reply, role: .model)) + if !keepRunning { continuation.finish() } + Task { await started.signal() } + } + } + + func cancel() { + cancelCount.withLock { $0 += 1 } + continuation.withLock { $0 }?.finish(throwing: CancellationError()) + Task { await cancelled.signal() } + } + } + + private final class LiteRTImageURLProtocol: URLProtocol { + override class func canInit(with request: URLRequest) -> Bool { true } + override class func canonicalRequest(for request: URLRequest) -> URLRequest { request } + + override func startLoading() { + let url = request.url! + if url.lastPathComponent == "network" { + client?.urlProtocol(self, didFailWithError: URLError(.notConnectedToInternet)) + return + } + let response = HTTPURLResponse( + url: url, + statusCode: url.lastPathComponent == "missing" ? 404 : 200, + httpVersion: "HTTP/1.1", + headerFields: nil + )! + client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed) + client?.urlProtocol(self, didLoad: Data("image bytes".utf8)) + client?.urlProtocolDidFinishLoading(self) + } + + override func stopLoading() {} + } +#endif From 7e146ce2b96df335413c155ab874842a29d9a99b Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Fri, 11 Sep 2026 05:33:29 -0700 Subject: [PATCH 5/6] Keep optional CI dependency resolution out of the lockfile --- Package.resolved | 198 ----------------------------------------------- 1 file changed, 198 deletions(-) diff --git a/Package.resolved b/Package.resolved index 5d72bfa5..ccb20bce 100644 --- a/Package.resolved +++ b/Package.resolved @@ -1,15 +1,6 @@ { "originHash" : "872262beb92df73ada611fd9ca3cc641833417deb1e27d47f5df8176d080243d", "pins" : [ - { - "identity" : "async-http-client", - "kind" : "remoteSourceControl", - "location" : "https://github.com/swift-server/async-http-client.git", - "state" : { - "revision" : "f95c908967e98c68c5ce3fd61a7974e7e869e303", - "version" : "1.36.1" - } - }, { "identity" : "eventsource", "kind" : "remoteSourceControl", @@ -46,24 +37,6 @@ "version" : "2.10549.0" } }, - { - "identity" : "mlx-swift", - "kind" : "remoteSourceControl", - "location" : "https://github.com/ml-explore/mlx-swift", - "state" : { - "revision" : "0bb916c67f4b9e5c682cbe02a42c701c93ab5021", - "version" : "0.31.6" - } - }, - { - "identity" : "mlx-swift-lm", - "kind" : "remoteSourceControl", - "location" : "https://github.com/ml-explore/mlx-swift-lm", - "state" : { - "revision" : "bd4b7434e6bdb588c7ef55706ff8904cb7fd4c57", - "version" : "3.31.4" - } - }, { "identity" : "partialjsondecoder", "kind" : "remoteSourceControl", @@ -73,24 +46,6 @@ "version" : "1.0.0" } }, - { - "identity" : "swift-algorithms", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-algorithms.git", - "state" : { - "revision" : "87e50f483c54e6efd60e885f7f5aa946cee68023", - "version" : "1.2.1" - } - }, - { - "identity" : "swift-argument-parser", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-argument-parser", - "state" : { - "revision" : "6a52f3251125d74daf04fcbd5e6f08a75d074382", - "version" : "1.8.2" - } - }, { "identity" : "swift-asn1", "kind" : "remoteSourceControl", @@ -100,15 +55,6 @@ "version" : "1.7.2" } }, - { - "identity" : "swift-async-algorithms", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-async-algorithms.git", - "state" : { - "revision" : "3da39bbc4e687d4192af7c9cf4eab805745a0b9c", - "version" : "1.1.5" - } - }, { "identity" : "swift-atomics", "kind" : "remoteSourceControl", @@ -118,15 +64,6 @@ "version" : "1.3.1" } }, - { - "identity" : "swift-certificates", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-certificates.git", - "state" : { - "revision" : "c8aece90ea05f9866bd392a5bf13b5cae56c0e03", - "version" : "1.20.0" - } - }, { "identity" : "swift-collections", "kind" : "remoteSourceControl", @@ -136,15 +73,6 @@ "version" : "1.6.0" } }, - { - "identity" : "swift-configuration", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-configuration.git", - "state" : { - "revision" : "be76c4ad929eb6c4bcaf3351799f2adf9e6848a9", - "version" : "1.2.0" - } - }, { "identity" : "swift-crypto", "kind" : "remoteSourceControl", @@ -154,33 +82,6 @@ "version" : "4.5.2" } }, - { - "identity" : "swift-distributed-tracing", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-distributed-tracing.git", - "state" : { - "revision" : "dc4030184203ffafbb2ec614352487235d747fe0", - "version" : "1.4.1" - } - }, - { - "identity" : "swift-http-structured-headers", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-http-structured-headers.git", - "state" : { - "revision" : "933538faa42c432d385f02e07df0ace7c5ecfc47", - "version" : "1.7.0" - } - }, - { - "identity" : "swift-http-types", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-http-types.git", - "state" : { - "revision" : "bff4b6903cdc99dda49649dd52f46c11cfd3ed50", - "version" : "1.8.0" - } - }, { "identity" : "swift-huggingface", "kind" : "remoteSourceControl", @@ -190,24 +91,6 @@ "version" : "0.10.1" } }, - { - "identity" : "swift-jinja", - "kind" : "remoteSourceControl", - "location" : "https://github.com/huggingface/swift-jinja.git", - "state" : { - "revision" : "12ea4955e380dae6b0cd98299effe48c7fe584fc", - "version" : "2.5.0" - } - }, - { - "identity" : "swift-log", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-log.git", - "state" : { - "revision" : "9c6fb14227f55d8f711ce3847dc2f419fb0ecacb", - "version" : "1.15.1" - } - }, { "identity" : "swift-nio", "kind" : "remoteSourceControl", @@ -217,69 +100,6 @@ "version" : "2.101.2" } }, - { - "identity" : "swift-nio-extras", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-nio-extras.git", - "state" : { - "revision" : "41449336c8ecfadac6b4b5be75f9c3c306e61ced", - "version" : "1.35.1" - } - }, - { - "identity" : "swift-nio-http2", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-nio-http2.git", - "state" : { - "revision" : "0f3e54e29c944c2e835ad52159da7d9e1c94ac69", - "version" : "1.46.0" - } - }, - { - "identity" : "swift-nio-ssl", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-nio-ssl.git", - "state" : { - "revision" : "03827c1a9fdb2b6b00a4e93ede8861520263af8c", - "version" : "2.37.4" - } - }, - { - "identity" : "swift-nio-transport-services", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-nio-transport-services.git", - "state" : { - "revision" : "67787bb645a5e67d2edcdfbe48a216cc549222d5", - "version" : "1.28.0" - } - }, - { - "identity" : "swift-numerics", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-numerics", - "state" : { - "revision" : "0c0290ff6b24942dadb83a929ffaaa1481df04a2", - "version" : "1.1.1" - } - }, - { - "identity" : "swift-service-context", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-service-context.git", - "state" : { - "revision" : "d0997351b0c7779017f88e7a93bc30a1878d7f29", - "version" : "1.3.0" - } - }, - { - "identity" : "swift-service-lifecycle", - "kind" : "remoteSourceControl", - "location" : "https://github.com/swift-server/swift-service-lifecycle", - "state" : { - "revision" : "7f9326b0326ff86e3646295ea6e891f68c471c5e", - "version" : "2.12.0" - } - }, { "identity" : "swift-syntax", "kind" : "remoteSourceControl", @@ -297,24 +117,6 @@ "revision" : "7502b711c92a17741fa625d722b0ccbd595d8ed1", "version" : "1.7.2" } - }, - { - "identity" : "swift-transformers", - "kind" : "remoteSourceControl", - "location" : "https://github.com/huggingface/swift-transformers", - "state" : { - "revision" : "c21fdcde390313a6d98d8e33a346f2c3486c3ab0", - "version" : "1.3.4" - } - }, - { - "identity" : "yyjson", - "kind" : "remoteSourceControl", - "location" : "https://github.com/ibireme/yyjson.git", - "state" : { - "revision" : "8b4a38dc994a110abaec8a400615567bd996105f", - "version" : "0.12.0" - } } ], "version" : 3 From 1e29d79a22f614ff41a5276e06ca927450ce78ee Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Fri, 11 Sep 2026 06:03:50 -0700 Subject: [PATCH 6/6] Fix LiteRT structured values and sampling options --- .../Models/LiteRTLanguageModel.swift | 136 +++++++++++++----- .../LiteRTLanguageModelBehaviorTests.swift | 106 +++++++++++++- 2 files changed, 204 insertions(+), 38 deletions(-) diff --git a/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift b/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift index 5e321e79..3cd7a6fd 100644 --- a/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift @@ -232,7 +232,11 @@ import Foundation ) } - let json = extractJSONObject(from: text) ?? text + guard let json = extractJSONValue(from: text) else { + throw LanguageModelSession.GenerationError.decodingFailure( + .init(debugDescription: "LiteRT did not generate a complete JSON value.") + ) + } let generatedContent = try GeneratedContent(json: json) let content = try type.init(generatedContent) return LanguageModelSession.Response( @@ -279,6 +283,7 @@ import Foundation ) var text = "" + var lastJSON: String? try await generateLiteRTResponse( conversation: conversation, prompt: plan.prompt, @@ -294,10 +299,12 @@ import Foundation rawContent: GeneratedContent(text) ) ) - } else if let json = extractJSONObject(from: text), + } else if let json = extractJSONValue(from: text, isFinal: false), + json != lastJSON, let raw = try? GeneratedContent(json: json), let parsed = try? type.init(raw) { + lastJSON = json continuation.yield( .init( content: parsed.asPartiallyGenerated(), @@ -310,6 +317,20 @@ import Foundation } } + if type != String.self { + guard let json = extractJSONValue(from: text) else { + throw LanguageModelSession.GenerationError.decodingFailure( + .init(debugDescription: "LiteRT did not generate a complete JSON value.") + ) + } + if json != lastJSON { + let raw = try GeneratedContent(json: json) + let parsed = try type.init(raw) + continuation.yield( + .init(content: parsed.asPartiallyGenerated(), rawContent: raw) + ) + } + } continuation.finish() } catch { continuation.finish(throwing: error) @@ -408,7 +429,7 @@ import Foundation if isTrigger, let schemaJSON, !schemaJSON.isEmpty { contents.append( .text( - "\n\nRespond with ONLY a JSON object that conforms to this JSON schema. " + "\n\nRespond with ONLY a JSON value that conforms to this JSON schema. " + "Output valid JSON and nothing else:\n\(schemaJSON)" ) ) @@ -494,33 +515,68 @@ import Foundation return String(data: data, encoding: .utf8) ?? "" } - /// Extracts the first balanced JSON object from model text - /// (strips surrounding prose and code fences). - private func extractJSONObject(from text: String) -> String? { - guard let start = text.firstIndex(of: "{") else { return nil } - var depth = 0 - var inString = false - var escaped = false - var index = start - while index < text.endIndex { - let character = text[index] - if inString { - if escaped { - escaped = false - } else if character == "\\" { - escaped = true - } else if character == "\"" { - inString = false + /// Extracts the first complete JSON value from model text, + /// stripping surrounding prose and code fences. + /// Scalars at the end of a stream must wait for the next delimiter or completion, + /// because another chunk can extend a number or literal. + private func extractJSONValue(from text: String, isFinal: Bool = true) -> String? { + var start = text.startIndex + while start < text.endIndex { + let first = text[start] + let isContainer = first == "{" || first == "[" + let isString = first == "\"" + let isBoundary = + start == text.startIndex + || !(text[text.index(before: start)].isLetter + || text[text.index(before: start)].isNumber + || text[text.index(before: start)] == ".") + guard + isContainer || isString + || (isBoundary && "-0123456789tfn".contains(first)) + else { + start = text.index(after: start) + continue + } + + var end = start + if isContainer || isString { + var depth = 0 + var inString = false + var escaped = false + repeat { + let character = text[end] + if inString { + if escaped { + escaped = false + } else if character == "\\" { + escaped = true + } else if character == "\"" { + inString = false + } + } else if character == "\"" { + inString = true + } else if character == "{" || character == "[" { + depth += 1 + } else if character == "}" || character == "]" { + depth -= 1 + } + end = text.index(after: end) + } while end < text.endIndex && (inString || depth > 0) + guard !inString, depth == 0 else { return nil } + } else { + while end < text.endIndex, + text[end].isLetter || text[end].isNumber || ".+-".contains(text[end]) + { + end = text.index(after: end) } - } else if character == "\"" { - inString = true - } else if character == "{" { - depth += 1 - } else if character == "}" { - depth -= 1 - if depth == 0 { return String(text[start ... index]) } + if end == text.endIndex && !isFinal { return nil } + } + + let candidate = String(text[start ..< end]) + if (try? JSONSerialization.jsonObject(with: Data(candidate.utf8), options: .fragmentsAllowed)) != nil { + return candidate } - index = text.index(after: index) + start = end } return nil } @@ -530,6 +586,7 @@ import Foundation private func makeSampler(for options: GenerationOptions, structured: Bool) -> SamplerConfig? { var topK = 40 var topP = 0.95 + var seed: UInt64 = 0 // Lower default temperature for structured / tool generation // (more reliable JSON). var temperature = structured ? 0.0 : 0.8 @@ -540,13 +597,25 @@ import Foundation switch sampling.mode { case .greedy: temperature = 0.0 - case .topK(let k, _): - topK = k - case .nucleus(let probabilityThreshold, _): + case .topK(let k, let randomSeed): + topK = min(k, Int(Int32.max)) + topP = 1.0 + seed = randomSeed ?? 0 + case .nucleus(let probabilityThreshold, let randomSeed): + // LiteRT clamps top-k to the vocabulary size. + topK = Int(Int32.max) topP = probabilityThreshold + seed = randomSeed ?? 0 } } - return try? SamplerConfig(topK: topK, topP: Float(topP), temperature: Float(temperature)) + // The Swift wrapper converts seeds to the runtime's signed 32-bit representation. + // Preserve the low 32 bits without trapping on larger UInt64 values. + return try? SamplerConfig( + topK: topK, + topP: Float(topP), + temperature: Float(temperature), + seed: Int(Int32(truncatingIfNeeded: seed)) + ) } // MARK: - Tool Calling @@ -574,7 +643,8 @@ import Foundation from text: String, tools: [any Tool] ) -> (name: String, arguments: String)? { - guard let json = extractJSONObject(from: text), + guard let start = text.firstIndex(of: "{"), + let json = extractJSONValue(from: String(text[start...])), let data = json.data(using: .utf8), let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any], let call = object["tool_call"] as? [String: Any], diff --git a/Tests/AnyLanguageModelTests/LiteRTLanguageModelBehaviorTests.swift b/Tests/AnyLanguageModelTests/LiteRTLanguageModelBehaviorTests.swift index d794108d..10b04654 100644 --- a/Tests/AnyLanguageModelTests/LiteRTLanguageModelBehaviorTests.swift +++ b/Tests/AnyLanguageModelTests/LiteRTLanguageModelBehaviorTests.swift @@ -8,6 +8,88 @@ import Testing @Suite("LiteRTLanguageModel behavior", .timeLimit(.minutes(1))) struct LiteRTLanguageModelBehaviorTests { + @Test(arguments: ["plain", "fenced", "prose"], [false, true]) + func generatesArraysAndScalars(wrapper: String, streaming: Bool) async throws { + func check(_ json: String, equals expected: Value) async throws { + let reply: String + switch wrapper { + case "fenced": reply = "```json\n\(json)\n```" + case "prose": reply = "The result is:\n\(json)\nThat is the answer." + default: reply = json + } + let conversation = StubLiteRTConversation(chunks: reply.map(String.init)) + let runtime = StubLiteRTRuntime([conversation]) + let session = LanguageModelSession(model: LiteRTLanguageModel(load: { runtime })) + if streaming { + var values: [Value] = [] + for try await snapshot in session.streamResponse(to: "Generate", generating: Value.self) { + values.append(try Value(snapshot.rawContent)) + } + #expect(values == [expected]) + } else { + #expect(try await session.respond(to: "Generate", generating: Value.self).content == expected) + } + #expect(conversation.requests.withLock { $0.first?.text.contains("JSON value") } == true) + } + try await check(#"{"answer":"hello"}"#, equals: LiteRTTestAnswer(answer: "hello")) + try await check("[]", equals: [Int]()) + try await check("[1,2,3]", equals: [1, 2, 3]) + try await check("[[1,2],[],[3]]", equals: [[1, 2], [], [3]]) + try await check("42", equals: 42) + try await check("-1.25e+2", equals: -125.0) + try await check("true", equals: true) + try await check("false", equals: false) + try await check(#"["a]b", "a\"b", "a\\b"]"#, equals: ["a]b", "a\"b", "a\\b"]) + } + + @Test(arguments: ["[1,", "[1:2]", "1e", "tru", "falsehood"], [false, true]) + func rejectsIncompleteOrInvalidStructuredResponses(reply: String, streaming: Bool) async throws { + let runtime = StubLiteRTRuntime([StubLiteRTConversation(reply: reply)]) + let session = LanguageModelSession(model: LiteRTLanguageModel(load: { runtime })) + await #expect { + if streaming { + for try await _ in session.streamResponse(to: "Generate", generating: [Int].self) { + Issue.record("Incomplete JSON must not produce a snapshot") + } + } else { + _ = try await session.respond(to: "Generate", generating: [Int].self) + } + } throws: { error in + guard case LanguageModelSession.GenerationError.decodingFailure = error else { return false } + return true + } + } + + @Test(arguments: [false, true], [UInt64(123), UInt64(Int32.max) + 1, UInt64.max]) + func preservesExplicitSamplingModesAndSeeds(streaming: Bool, seed: UInt64) async throws { + for sampling in [ + GenerationOptions.SamplingMode.random(top: 10, seed: seed), + .random(probabilityThreshold: 0.8, seed: seed), + ] { + let runtime = StubLiteRTRuntime([StubLiteRTConversation(reply: "Hello")]) + let session = LanguageModelSession(model: LiteRTLanguageModel(load: { runtime })) + let options = GenerationOptions(sampling: sampling, temperature: 0.7) + if streaming { + for try await _ in session.streamResponse(to: "Hello", options: options) {} + } else { + _ = try await session.respond(to: "Hello", options: options) + } + let config = try #require(runtime.configurations.withLock { $0.first?.sampler }) + switch sampling.mode { + case .topK: + #expect(config.topK == 10) + #expect(config.topP == 1.0) + case .nucleus: + #expect(config.topK == Int(Int32.max)) + #expect(config.topP == 0.8) + case .greedy: Issue.record("Unexpected sampling mode") + } + #expect(config.temperature == 0.7) + let expectedSeed: Int = seed == 123 ? 123 : (seed == UInt64.max ? -1 : Int(Int32.min)) + #expect(config.seed == expectedSeed) + } + } + @Test(arguments: [false, true]) func forwardsResponseTokenLimit(streaming: Bool) async throws { let conversation = StubLiteRTConversation(reply: "Hello") @@ -229,7 +311,7 @@ import Testing private let toolCall = #"{"tool_call":{"name":"lookup","arguments":{"query":"answer"}}}"# @Generable - private struct LiteRTTestAnswer { + private struct LiteRTTestAnswer: Equatable { var answer: String } @@ -275,7 +357,14 @@ import Testing } private final class StubLiteRTRuntime: LiteRTRuntime { + struct Sampler: Sendable { + var topK: Int + var topP: Float + var temperature: Float + var seed: Int + } struct Configuration: Sendable { + var sampler: Sampler? var systemMessage: String? var history: [String] } @@ -291,6 +380,9 @@ import Testing let index = records.count records.append( Configuration( + sampler: config.samplerConfig.map { + Sampler(topK: $0.topK, topP: $0.topP, temperature: $0.temperature, seed: $0.seed) + }, systemMessage: config.systemMessage?.toString, history: config.initialMessages.map(\.toString) ) @@ -308,7 +400,7 @@ import Testing var images: [Data] var maxOutputTokens: Int? } - let reply: String + let chunks: [String] let keepRunning: Bool let requests = Locked<[Request]>([]) let cancelCount = Locked(0) @@ -316,8 +408,12 @@ import Testing let started = LiteRTTestSignal() let cancelled = LiteRTTestSignal() - init(reply: String, keepRunning: Bool = false) { - self.reply = reply + convenience init(reply: String, keepRunning: Bool = false) { + self.init(chunks: [reply], keepRunning: keepRunning) + } + + init(chunks: [String], keepRunning: Bool = false) { + self.chunks = chunks self.keepRunning = keepRunning } @@ -335,7 +431,7 @@ import Testing } return AsyncThrowingStream { continuation in self.continuation.withLock { $0 = continuation } - continuation.yield(Message(reply, role: .model)) + for chunk in chunks { continuation.yield(Message(chunk, role: .model)) } if !keepRunning { continuation.finish() } Task { await started.signal() } }