diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 25f77fbd..82cfaa96 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -10,11 +10,6 @@ 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) || '' }}
@@ -68,10 +63,10 @@ jobs:
run: swift format lint --strict --recursive .
- name: Build
- run: swift build --build-tests --traits MLX,Llama,CoreML,LiteRT${{ matrix.traits != '' && format(',{0}', matrix.traits) || '' }}
+ run: swift build --build-tests --traits MLX,Llama,CoreML${{ matrix.traits != '' && format(',{0}', matrix.traits) || '' }}
- name: Test
- run: swift test --skip-build --traits MLX,Llama,CoreML,LiteRT${{ matrix.traits != '' && format(',{0}', matrix.traits) || '' }}
+ run: swift test --skip-build --traits MLX,Llama,CoreML${{ 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 0e9e576f..e486dbf2 100644
--- a/Package.resolved
+++ b/Package.resolved
@@ -1,5 +1,5 @@
{
- "originHash" : "6cd871e589ddd4a6cf4040c379b16a902babca7ad3cf69dde3383238c658af31",
+ "originHash" : "b46e46d156bb5bfeea0c5c4e9b59d51ebea883a40627ee34aed14b97f142276d",
"pins" : [
{
"identity" : "eventsource",
@@ -19,15 +19,6 @@
"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" : "llama.swift",
"kind" : "remoteSourceControl",
diff --git a/Package.swift b/Package.swift
index 0d4a661a..2c4a8d63 100644
--- a/Package.swift
+++ b/Package.swift
@@ -25,7 +25,6 @@ let package = Package(
.trait(name: "CoreML"),
.trait(name: "MLX"),
.trait(name: "Llama"),
- .trait(name: "LiteRT"),
.trait(name: "AsyncHTTPClient"),
.default(enabledTraits: []),
],
@@ -44,7 +43,6 @@ let package = Package(
.package(url: "https://github.com/mattt/llama.swift", .upToNextMajor(from: "2.10549.0")),
.package(url: "https://github.com/mattt/PartialJSONDecoder", from: "1.0.0"),
.package(url: "https://github.com/ml-explore/mlx-swift-lm", from: "3.31.4"),
- .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 +77,7 @@ let package = Package(
.product(
name: "HuggingFace",
package: "swift-huggingface",
- condition: .when(traits: ["MLX", "LiteRT"])
+ condition: .when(traits: ["MLX"])
),
.product(
name: "Tokenizers",
@@ -96,11 +94,6 @@ let package = Package(
package: "llama.swift",
condition: .when(traits: ["Llama"])
),
- .product(
- name: "LiteRTLM",
- package: "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 38c2eac9..ac5f7836 100644
--- a/README.md
+++ b/README.md
@@ -73,8 +73,6 @@ 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] 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)
@@ -100,7 +98,7 @@ Add this package to your `Package.swift`:
```swift
dependencies: [
- .package(url: "https://github.com/huggingface/AnyLanguageModel", from: "0.10.0")
+ .package(url: "https://github.com/huggingface/AnyLanguageModel", from: "0.11.0")
]
```
@@ -119,9 +117,6 @@ 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 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:
@@ -131,7 +126,7 @@ To enable specific traits, specify them in your package's dependencies:
dependencies: [
.package(
url: "https://github.com/huggingface/AnyLanguageModel.git",
- from: "0.10.0",
+ from: "0.11.0",
traits: ["CoreML", "MLX"] // Enable CoreML and MLX support
)
]
@@ -148,7 +143,7 @@ dependencies: [
> dependencies: [
> .package(
> url: "https://github.com/huggingface/AnyLanguageModel.git",
-> from: "0.10.0",
+> from: "0.11.0",
> traits: ["CoreML", "MLX", "Llama"]
> ),
> .package(url: "https://github.com/huggingface/swift-transformers", from: "1.0.0"), // CoreML
@@ -237,94 +232,6 @@ Your app can now import `AnyLanguageModel` with MLX support enabled.
> For a working example of package traits in an Xcode app project,
> see [chat-ui-swift](https://github.com/mattt/chat-ui-swift).
-## Troubleshooting
-
-### LiteRT-LM Checkout Fails with a Git LFS Smudge Error
-
-If Git LFS is installed and configured globally,
-a fresh dependency checkout may fail
-with this error from LiteRT-LM 0.17.0:
-
-```text
-error: 'litert-lm': Couldn’t check out revision ‘e9fd8c53ff968071774206163027dd84bedfe925’:
- Downloading prebuilt/android_arm64/libGemmaModelConstraintProvider.so (20 MB)
- Error downloading object: prebuilt/android_arm64/libGemmaModelConstraintProvider.so (2db0cfa): Smudge error: Error downloading prebuilt/android_arm64/libGemmaModelConstraintProvider.so (2db0cfa5d45391df18e6c8de4b1e5ffbe1882d3695c0a0b5e087e4e482e1d680): error transferring "2db0cfa5d45391df18e6c8de4b1e5ffbe1882d3695c0a0b5e087e4e482e1d680": [0] remote missing object 2db0cfa5d45391df18e6c8de4b1e5ffbe1882d3695c0a0b5e087e4e482e1d680
-```
-
-Git LFS tries to download a prebuilt Android library during checkout,
-but the referenced object is missing from the LFS remote used by that checkout.
-Swift Package Manager creates dependency checkouts from a local repository mirror,
-which Git LFS can treat as its remote
-instead of the original GitHub repository.
-That mirror may lack LFS objects required by the selected version,
-even when those objects are available on GitHub.
-See [upstream issue #2407](https://github.com/google-ai-edge/LiteRT-LM/issues/2407)
-for discussion and updates,
-and [upstream PR #3563](https://github.com/google-ai-edge/LiteRT-LM/pull/3563)
-for a proposed fix that explicitly configures the GitHub LFS endpoint.
-This can happen even when the `LiteRT` trait is disabled,
-because Swift Package Manager still resolves the package dependency.
-
-The failure occurs during dependency checkout,
-before compilation.
-It can affect `swift build`, `swift test`, `xcodebuild`,
-and package resolution in Xcode,
-blocking Build or Run.
-An existing successful checkout may hide the problem
-until dependencies are fetched again,
-for example after deleting `.build`
-or resetting Xcode's package caches.
-It doesn't affect an already-built app at runtime.
-
-To skip LFS downloads for a Swift Package Manager command,
-prefix it with `GIT_LFS_SKIP_SMUDGE=1`:
-
-```bash
-GIT_LFS_SKIP_SMUDGE=1 swift build
-GIT_LFS_SKIP_SMUDGE=1 swift test
-```
-
-For a clean build, use:
-
-```bash
-rm -rf .build && GIT_LFS_SKIP_SMUDGE=1 swift test
-```
-
-For `xcodebuild`,
-use the same prefix
-with your usual project or workspace and scheme arguments.
-For example,
-replacing `MyApp` with your project and scheme names:
-
-```bash
-GIT_LFS_SKIP_SMUDGE=1 xcodebuild -project MyApp.xcodeproj -scheme MyApp build
-```
-
-For Xcode's Build and Run actions,
-quit Xcode completely,
-then launch its executable from Terminal with the variable set:
-
-```bash
-GIT_LFS_SKIP_SMUDGE=1 /Applications/Xcode.app/Contents/MacOS/Xcode
-```
-
-Adjust the path
-if Xcode is installed under a different name or location,
-then open your project and retry package resolution or the build.
-Setting the variable in a terminal
-doesn't affect an already-running Xcode.
-Adding it to a scheme's Run environment variables won't help either:
-those variables apply to the launched app,
-after package resolution.
-
-LiteRT-LM's [Swift package manifest](https://github.com/google-ai-edge/LiteRT-LM/blob/e9fd8c53ff968071774206163027dd84bedfe925/Package.swift)
-downloads Apple XCFrameworks separately from release assets,
-so its Git LFS binaries aren't needed for Swift Package Manager builds.
-The repository's [CI workflow](.github/workflows/ci.yml)
-already uses this workaround.
-The environment variable applies only to the command and its subprocesses;
-it doesn't change your global Git LFS configuration.
-
## API Credentials and Security
When using third-party language model providers like OpenAI, Anthropic, or Google Gemini,
@@ -459,14 +366,13 @@ 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, LiteRT-LM, and Ollama,
+For MLX and Ollama,
use a vision-capable model
(for example, a VLM or `-vl` variant).
@@ -563,7 +469,7 @@ Enable the trait in Package.swift:
```swift
.package(
url: "https://github.com/huggingface/AnyLanguageModel.git",
- from: "0.10.0",
+ from: "0.11.0",
traits: ["CoreML"]
)
```
@@ -649,7 +555,7 @@ Enable the trait in Package.swift:
```swift
.package(
url: "https://github.com/huggingface/AnyLanguageModel.git",
- from: "0.10.0",
+ from: "0.11.0",
traits: ["MLX"]
)
```
@@ -673,7 +579,7 @@ Enable the trait in Package.swift:
```swift
.package(
url: "https://github.com/huggingface/AnyLanguageModel.git",
- from: "0.10.0",
+ from: "0.11.0",
traits: ["Llama"]
)
```
@@ -704,38 +610,15 @@ 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(modelFileURL: modelURL)
-
-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:
+
-```swift
-// Any Hugging Face repo
-let model = LiteRTLanguageModel(
- huggingFaceRepo: "litert-community/gemma-4-E4B-it-litert-lm",
- fileName: "gemma-4-E4B-it.litertlm")
-```
+### LiteRT-LM
-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).
+LiteRT-LM support was added in 0.10.0 and removed in 0.11.0
+because its dependency could prevent builds even when the backend was disabled.
+See [upstream issue #2407](https://github.com/google-ai-edge/LiteRT-LM/issues/2407)
+for details.
+The `LiteRT` trait and `LiteRTLanguageModel` are no longer available.
### Ollama
diff --git a/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift b/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift
deleted file mode 100644
index 3cd7a6fd..00000000
--- a/Sources/AnyLanguageModel/Models/LiteRTLanguageModel.swift
+++ /dev/null
@@ -1,715 +0,0 @@
-import Foundation
-
-#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.
- ///
- /// 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)
- /// let session = LanguageModelSession(model: model)
- /// let response = try await session.respond(to: "What is the capital of France?")
- /// ```
- ///
- /// 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 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: LiteRTModelLoader
- private let imageSession: URLSession
-
- /// Creates a model from a local `.litertlm` file.
- ///
- /// - 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.
- /// - maxTokens: Context (KV cache) budget.
- public init(
- modelFileURL: URL,
- backend: Backend = .gpu,
- visionBackend: Backend? = nil,
- audioBackend: Backend? = nil,
- maxTokens: Int = 2048
- ) {
- self.imageSession = .shared
- self.engine = LiteRTModelLoader {
- guard modelFileURL.isFileURL,
- FileManager.default.fileExists(atPath: modelFileURL.path)
- else {
- throw CocoaError(.fileReadNoSuchFile, userInfo: [NSURLErrorKey: modelFileURL])
- }
- return try await makeEngine(
- modelPath: modelFileURL.path,
- backend: backend,
- visionBackend: visionBackend,
- audioBackend: audioBackend,
- maxTokens: maxTokens
- )
- }
- }
-
- /// 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: 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.
- /// - maxTokens: Context (KV cache) budget.
- /// - 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",
- backend: Backend = .gpu,
- visionBackend: Backend? = nil,
- audioBackend: Backend? = nil,
- maxTokens: Int = 2048,
- hub: HubClient? = nil,
- downloadProgress: Progress? = nil
- ) {
- self.imageSession = .shared
- self.engine = LiteRTModelLoader {
- guard let repo = Repo.ID(rawValue: huggingFaceRepo) else {
- throw URLError(.badURL)
- }
- let destination = try await (hub ?? .default).downloadFile(
- at: fileName,
- from: repo,
- revision: revision,
- progress: downloadProgress
- )
- return try await makeEngine(
- modelPath: destination.path,
- backend: backend,
- visionBackend: visionBackend,
- audioBackend: audioBackend,
- maxTokens: maxTokens
- )
- }
- }
-
- 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?
- ) {
- 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 = try await makePlan(
- from: session.transcript,
- fallbackPrompt: prompt.description,
- schemaJSON: includeSchemaInPrompt ? schemaJSON : nil,
- tools: tools,
- imageSession: imageSession
- )
- 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.makeConversation(
- config: ConversationConfig(
- systemMessage: plan.systemMessage,
- initialMessages: plan.history,
- samplerConfig: sampler
- )
- )
-
- text = ""
- try await generateLiteRTResponse(
- conversation: conversation,
- prompt: plan.prompt,
- maximumResponseTokens: options.maximumResponseTokens
- ) { chunk in
- text += chunk
- }
-
- guard !tools.isEmpty,
- 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(
- name: parsed.name,
- argumentsJSON: parsed.arguments,
- session: session
- )
- 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,
- 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)
- )
- }
-
- 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(
- 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 = try await makePlan(
- from: session.transcript,
- fallbackPrompt: prompt.description,
- schemaJSON: includeSchemaInPrompt ? schemaJSON : nil,
- tools: [],
- imageSession: imageSession
- )
- let conversation = try await engine.makeConversation(
- config: ConversationConfig(
- systemMessage: plan.systemMessage,
- initialMessages: plan.history,
- samplerConfig: makeSampler(for: options, structured: schemaJSON != nil)
- )
- )
-
- var text = ""
- var lastJSON: String?
- try await generateLiteRTResponse(
- conversation: conversation,
- prompt: plan.prompt,
- maximumResponseTokens: options.maximumResponseTokens
- ) { delta in
- guard !delta.isEmpty else { return }
- text += delta
-
- if type == String.self {
- continuation.yield(
- .init(
- content: (text as! Content).asPartiallyGenerated(),
- rawContent: GeneratedContent(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(),
- rawContent: raw
- )
- )
- } else {
- // Structured responses stream as incomplete JSON fragments.
- // Skip snapshots until the accumulated JSON parses cleanly.
- }
- }
-
- 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)
- }
- }
-
- continuation.onTermination = { _ in
- task.cancel()
- }
- }
-
- return LanguageModelSession.ResponseStream(stream: stream)
- }
- }
-
- // MARK: - Engine Bring-Up
-
- private func makeEngine(
- modelPath: String,
- backend: Backend,
- visionBackend: Backend?,
- audioBackend: Backend?,
- maxTokens: Int
- ) async throws -> Engine {
- let caches = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask).first
- let config = try EngineConfig(
- modelPath: modelPath,
- backend: backend,
- visionBackend: visionBackend,
- audioBackend: audioBackend,
- maxNumTokens: maxTokens,
- cacheDir: caches?.path
- )
- 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],
- imageSession: URLSession
- ) async throws -> GenerationPlan {
- let entries = Array(transcript)
- let triggerIndex = entries.lastIndex { entry in
- switch entry {
- case .prompt, .toolOutput: return true
- default: return false
- }
- }
-
- var systemText: [String] = []
- let describedTools = tools.filter(\.includesSchemaInInstructions)
- if !describedTools.isEmpty {
- systemText.append(toolInstructions(describedTools))
- }
- 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 = try await messageContents(of: prompt.segments, session: imageSession)
- if isTrigger, let schemaJSON, !schemaJSON.isEmpty {
- contents.append(
- .text(
- "\n\nRespond with ONLY a JSON value 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],
- session: URLSession
- ) async throws -> [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 {
- 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))
- }
- }
- }
- }
- return contents.isEmpty ? [.text("")] : contents
- }
-
- /// 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
- switch segment {
- case .text(let text): return text.content
- case .structure(let structure): return structure.content.jsonString
- case .image: 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 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)
- }
- 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
- }
- start = end
- }
- return nil
- }
-
- // MARK: - Sampling
-
- 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
- if let explicit = options.temperature {
- temperature = explicit
- }
- if let sampling = options.sampling {
- switch sampling.mode {
- case .greedy:
- temperature = 0.0
- 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
- }
- }
- // 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
-
- 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 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],
- 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/Sources/AnyLanguageModel/Models/LiteRTRuntime.swift b/Sources/AnyLanguageModel/Models/LiteRTRuntime.swift
deleted file mode 100644
index 401f0dea..00000000
--- a/Sources/AnyLanguageModel/Models/LiteRTRuntime.swift
+++ /dev/null
@@ -1,131 +0,0 @@
-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
deleted file mode 100644
index 10b04654..00000000
--- a/Tests/AnyLanguageModelTests/LiteRTLanguageModelBehaviorTests.swift
+++ /dev/null
@@ -1,470 +0,0 @@
-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: ["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")
- 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: Equatable {
- 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 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]
- }
- 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(
- 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)
- )
- )
- 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 chunks: [String]
- let keepRunning: Bool
- let requests = Locked<[Request]>([])
- let cancelCount = Locked(0)
- let continuation = Locked.Continuation?>(nil)
- let started = LiteRTTestSignal()
- let cancelled = LiteRTTestSignal()
-
- 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
- }
-
- 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 }
- for chunk in chunks { continuation.yield(Message(chunk, 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
diff --git a/Tests/AnyLanguageModelTests/LiteRTLanguageModelTests.swift b/Tests/AnyLanguageModelTests/LiteRTLanguageModelTests.swift
deleted file mode 100644
index 6c45cdda..00000000
--- a/Tests/AnyLanguageModelTests/LiteRTLanguageModelTests.swift
+++ /dev/null
@@ -1,96 +0,0 @@
-import Foundation
-import Testing
-
-@testable import AnyLanguageModel
-
-#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 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!)
- )
- }
-
- @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