Merge pull request #64 from rryam/feature/mlx-package-traits

Add MLX Package Trait to reduce dependency footprint

Author: rudrankriyam

Package.resolved

@@ -1,4 +1,4 @@
-// swift-tools-version: 6.0
+// swift-tools-version: 6.1
 // The swift-tools-version declares the minimum version of Swift required to build this package.
 
 import PackageDescription
@@ -34,9 +34,18 @@ let package = Package(
       targets: ["VecturaMLXCLI"]
     ),
   ],
+  traits: [
+    .trait(
+      name: "MLX",
+      description: "Enable MLX-based embeddings for GPU-accelerated inference"
+    ),
+  ],
   dependencies: [
+    // Always included - lightweight dependencies
     .package(url: "https://github.com/jkrukowski/swift-embeddings.git", from: "0.0.21"),
     .package(url: "https://github.com/apple/swift-argument-parser.git", from: "1.4.0"),
+
+    // MLX dependency - only loaded when MLX trait is enabled via target conditions
     .package(url: "https://github.com/ml-explore/mlx-swift-lm/", from: "2.30.3"),
   ],
   targets: [
@@ -54,7 +63,11 @@ let package = Package(
       name: "VecturaMLXKit",
       dependencies: [
         "VecturaKit",
-        .product(name: "MLXEmbedders", package: "mlx-swift-lm"),
+        .product(
+          name: "MLXEmbedders",
+          package: "mlx-swift-lm",
+          condition: .when(traits: ["MLX"])
+        ),
       ]
     ),
     .target(
@@ -77,7 +90,7 @@ let package = Package(
       name: "VecturaMLXCLI",
       dependencies: [
         "VecturaKit",
-        "VecturaMLXKit",
+        .target(name: "VecturaMLXKit", condition: .when(traits: ["MLX"])),
         .product(name: "ArgumentParser", package: "swift-argument-parser"),
       ]
     ),
@@ -87,7 +100,9 @@ let package = Package(
     ),
     .executableTarget(
       name: "TestMLXExamples",
-      dependencies: ["VecturaMLXKit"]
+      dependencies: [
+        .target(name: "VecturaMLXKit", condition: .when(traits: ["MLX"]))
+      ]
     ),
     .executableTarget(
       name: "TestNLExamples",
@@ -99,7 +114,9 @@ let package = Package(
     ),
     .testTarget(
       name: "VecturaMLXKitTests",
-      dependencies: ["VecturaMLXKit"]
+      dependencies: [
+        .target(name: "VecturaMLXKit", condition: .when(traits: ["MLX"]))
+      ]
     ),
     .testTarget(
       name: "VecturaNLKitTests",

Package.swift

@@ -10,7 +10,7 @@ The framework offers `VecturaKit` as the core vector database with pluggable emb
 It also includes CLI tools (`vectura-cli` and `vectura-mlx-cli`) for easily trying out the package.
 
 <p align="center">
-  <img src="https://img.shields.io/badge/Swift-6.0+-fa7343?style=flat&logo=swift&logoColor=white" alt="Swift 6.0+">
+  <img src="https://img.shields.io/badge/Swift-6.1+-fa7343?style=flat&logo=swift&logoColor=white" alt="Swift 6.1+">
   <br>
   <img src="https://img.shields.io/badge/iOS-17.0+-000000?style=flat&logo=apple&logoColor=white" alt="iOS 17.0+">
   <img src="https://img.shields.io/badge/macOS-14.0+-000000?style=flat&logo=apple&logoColor=white" alt="macOS 14.0+">
@@ -40,6 +40,7 @@ Explore the following books to understand more about AI and iOS development:
 - [Supported Platforms](#supported-platforms)
 - [Installation](#installation)
   - [Swift Package Manager](#swift-package-manager)
+  - [Package Traits (MLX Support)](#package-traits-mlx-support)
   - [Dependencies](#dependencies)
 - [Usage](#usage)
   - [Import VecturaKit](#import-vecturakit)
@@ -102,10 +103,32 @@ To integrate VecturaKit into your project using Swift Package Manager, add the f
 
 ```swift
 dependencies: [
-    .package(url: "https://github.com/rryam/VecturaKit.git", from: "2.5.2"),
+    .package(url: "https://github.com/rryam/VecturaKit.git", from: "3.0.0"),
 ],
 ```
 
+### Package Traits (MLX Support)
+
+Starting with version 3.0.0, VecturaKit uses [Package Traits](https://github.com/swiftlang/swift-evolution/blob/main/proposals/0450-swiftpm-package-traits.md) (SE-0450) to reduce dependency footprint. The heavy MLX dependencies are only downloaded when explicitly enabled.
+
+**For VecturaKit or VecturaNLKit only (lightweight):**
+
+```swift
+dependencies: [
+    .package(url: "https://github.com/rryam/VecturaKit.git", from: "3.0.0"),
+],
+```
+
+**For VecturaMLXKit (enables MLX dependencies):**
+
+```swift
+dependencies: [
+    .package(url: "https://github.com/rryam/VecturaKit.git", from: "3.0.0", traits: ["MLX"]),
+],
+```
+
+> **Note:** Package Traits require Swift 6.1+ (Xcode 16.3+). If using Xcode's "Add Package Dependencies" UI for MLX support, you'll need to manually edit your `Package.swift` to add the `traits: ["MLX"]` parameter.
+
 ### Dependencies
 
 VecturaKit uses the following Swift packages:
@@ -418,6 +441,8 @@ let results = try await vectorDB.search(query: "search query")
 
 VecturaKit supports Apple's MLX framework through the `MLXEmbedder` for accelerated on-device machine learning performance.
 
+> **Important:** To use `VecturaMLXKit`, you must enable the `MLX` trait in your package dependency. See [Package Traits](#package-traits-mlx-support) for details.
+
 ### Import MLX Support
 
 ```swift
@@ -647,7 +672,7 @@ Common options for `vectura-cli`:
 -   `--dimension, -v`: Vector dimension (auto-detected by default)
 -   `--threshold, -t`: Minimum similarity threshold (default: 0.7)
 -   `--num-results, -n`: Number of results to return (default: 10)
--   `--model-id, -m`: Model ID for embeddings (default: "minishlab/potion-retrieval-32M")
+-   `--model-id, -m`: Model ID for embeddings (default: "minishlab/potion-base-4M")
 
 ### MLX CLI Tool (`vectura-mlx-cli`)
 
@@ -696,7 +721,7 @@ Contributions are welcome! Please fork the repository and submit a pull request
 ### Code Style
 
 - Follow SwiftLint rules (run `swiftlint lint`)
-- Use Swift 6.0+ features where appropriate
+- Use Swift 6.1+ features where appropriate
 - Maintain backward compatibility when possible
 - Document public APIs with DocC comments
 

README.md

@@ -1,5 +1,7 @@
 // Test script for VecturaMLXKit README examples
 import Foundation
+
+#if canImport(MLXEmbedders)
 import VecturaKit
 import VecturaMLXKit
 import MLXEmbedders
@@ -87,3 +89,12 @@ struct TestMLXExamples {
     debugPrint("Document count after reset: \(try await vectorDB.documentCount)")
   }
 }
+#else
+@main
+struct TestMLXExamples {
+  static func main() {
+    print("VecturaMLXKit is not available. Enable the MLX trait to run this example.")
+    print("Usage: swift run --traits MLX TestMLXExamples")
+  }
+}
+#endif

Sources/TestMLXExamples/TestMLXExamples.swift

@@ -100,7 +100,7 @@ extension VecturaCLI {
     var numResults: Int = 5
 
     @Option(name: [.long, .customShort("m")], help: "Model ID for embeddings")
-    var modelId: String = "minishlab/potion-retrieval-32M"
+    var modelId: String = "minishlab/potion-base-4M"
 
     mutating func run() async throws {
       print("VecturaKit Mock Demonstration")

Sources/VecturaCLI/VecturaCLI.swift

@@ -29,7 +29,7 @@ public enum VecturaModelSource: Sendable {
 public extension VecturaModelSource {
 
   /// The default model identifier when not otherwise specified.
-  static let defaultModelId = "minishlab/potion-retrieval-32M"
+  static let defaultModelId = "minishlab/potion-base-4M"
 
   /// The default model when not otherwise specified.
   static let `default` = Self.id(Self.defaultModelId)

Sources/VecturaKit/Embedder/VecturaModelSource.swift

@@ -1,5 +1,7 @@
 import ArgumentParser
 import Foundation
+
+#if canImport(MLXEmbedders)
 import MLXEmbedders
 import VecturaKit
 import VecturaMLXKit
@@ -280,3 +282,17 @@ extension VecturaMLXCLI {
     }
   }
 }
+#else
+@main
+struct VecturaMLXCLI: ParsableCommand {
+  static let configuration = CommandConfiguration(
+    commandName: "vectura-mlx",
+    abstract: "A CLI tool for VecturaKit vector database using MLX embeddings"
+  )
+
+  mutating func run() throws {
+    print("VecturaMLXKit is not available. Enable the MLX trait to use this CLI.")
+    print("Usage: swift run --traits MLX vectura-mlx-cli <command>")
+  }
+}
+#endif

Sources/VecturaMLXCLI/VecturaMLXCLI.swift

@@ -1,3 +1,4 @@
+#if canImport(MLXEmbedders)
 import Foundation
 import MLX
 import MLXEmbedders
@@ -126,3 +127,4 @@ enum EmbeddingError: Error {
   case unsupportedPoolingShape([Int])
   case vectorCountMismatch(expected: Int, received: Int)
 }
+#endif

Sources/VecturaMLXKit/MLXEmbedder.swift

@@ -1,3 +1,4 @@
+#if canImport(MLXEmbedders)
 import Foundation
 import Metal
 import MLX
@@ -11,6 +12,7 @@ import Testing
 /// 1. Metal device (GPU) availability
 /// 2. Metal Toolchain (install via: xcodebuild -downloadComponent MetalToolchain)
 /// 3. MLX device libraries to be available
+/// 4. MLX trait enabled (swift test --traits MLX)
 ///
 /// Run tests with: xcodebuild test -scheme VecturaMLXKitTests -destination 'platform=macOS'
 /// (swift test may not work due to Metal library compilation requirements)
@@ -238,3 +240,4 @@ struct VecturaMLXKitTests {
     #expect(results.isEmpty, "Search should return no results when the query does not match any document.")
   }
 }
+#endif

Tests/VecturaMLXKitTests/VecturaMLXKitTests.swift

@@ -28,14 +28,27 @@ definitions:
         # Set environment variables for model downloads
         export HF_HUB_OFFLINE=0
 
-        # Run core tests (excluding 30-min PerformanceTests suite)
+        # Run core tests (excluding PerformanceTests and MLX-related tests)
+        # MLX tests require the MLX trait to be enabled
         # Using --no-parallel to avoid concurrent model download issues
-        swift test --no-parallel --skip PerformanceTests
+        swift test --no-parallel \
+          --skip PerformanceTests \
+          --skip VecturaMLXKitTests
+    - &run_tests_with_mlx
+      name: Run Tests (with MLX)
+      script: |
+        #!/bin/zsh
+
+        # Set environment variables for model downloads
+        export HF_HUB_OFFLINE=0
+
+        # Run all tests with MLX trait enabled
+        swift test --no-parallel --traits MLX --skip PerformanceTests
 workflows:
   vecturakit:
     name: VecturaKit Workflow
     environment:
-      xcode: 26.0
+      xcode: 26.2
       vars:
         XCODE_SCHEME: "VecturaKit"
         APP_ID: "VecturaKit"
@@ -53,7 +66,7 @@ workflows:
   vecturamlxkit:
     name: VecturaMLXKit Workflow
     environment:
-      xcode: 26.0
+      xcode: 26.2
       vars:
         XCODE_SCHEME: "VecturaMLXKit"
         APP_ID: "VecturaMLXKit"
@@ -72,4 +85,4 @@ workflows:
 
           xcodebuild -downloadComponent MetalToolchain
       - *build_framework
-      - *run_tests
+      - *run_tests_with_mlx