Polish a few APIs, method names, and remove unnecessarily vending public extension (#122)

Co-authored-by: Tao Jia <tjia1818@users.noreply.github.com>

Author: tjia1818

swift/Sources/CoreAILanguageModels/DecodingStrategies/VanillaDecodingStrategy.swift

@@ -164,7 +164,7 @@ extension VanillaDecodingStrategy.VanillaDecodedSequence {
             if !flushed {
                 let count = generatedTokenCount
                 // Metrics recording is async; fire-and-forget on the abandon path.
-                Task { await PerformanceMetrics.shared.setGeneratedTokenCount(count) }
+                Task { await PerformanceMetrics.shared.recordGeneratedTokens(count) }
             }
         }
 
@@ -283,7 +283,7 @@ extension VanillaDecodingStrategy.VanillaDecodedSequence {
                 trailing = nil
             }
 
-            await PerformanceMetrics.shared.setGeneratedTokenCount(generatedTokenCount)
+            await PerformanceMetrics.shared.recordGeneratedTokens(generatedTokenCount)
             return trailing
         }
     }

swift/Sources/CoreAILanguageModels/LanguageModel/CoreAIRunner.swift

@@ -32,15 +32,15 @@ public struct CoreAIRunner {
         kvCacheStrategy: KVCacheStrategy = .auto
     ) throws {
         self.init(
-            from: try LanguageBundle(at: url),
+            bundle: try LanguageBundle(at: url),
             variant: variant,
             kvCacheStrategy: kvCacheStrategy
         )
     }
 
     /// Creates a runner from a LanguageBundle.
     public init(
-        from bundle: LanguageBundle,
+        bundle: LanguageBundle,
         variant: String? = nil,
         kvCacheStrategy: KVCacheStrategy = .auto
     ) {

swift/Sources/CoreAILanguageModels/Profiling/PerformanceMetrics.swift

@@ -28,9 +28,8 @@ public final class PerformanceMetrics {
     private var startInstant: ContinuousClock.Instant?
     private var endInstant: ContinuousClock.Instant?
 
-    private var promptTokenCount: Int = 0
-    private var generatedTokenCount: Int = 0
-    private var totalTokenCount: Int = 0
+    public private(set) var promptTokenCount: Int = 0
+    public private(set) var generatedTokenCount: Int = 0
 
     /// The shared PerformanceMetrics instance for production use.
     public static let shared = PerformanceMetrics()
@@ -61,17 +60,19 @@ public final class PerformanceMetrics {
 
     // MARK: - Token Counting
 
-    public func setPromptTokenCount(_ count: Int) {
-        promptTokenCount = count
+    /// Total number of prompt and generated tokens.
+    public var totalTokenCount: Int {
+        promptTokenCount + generatedTokenCount
     }
 
-    public func setGeneratedTokenCount(_ count: Int) {
-        generatedTokenCount = count
-        totalTokenCount = promptTokenCount + generatedTokenCount
+    /// Records the number of tokens in the prompt.
+    public func recordPromptTokens(_ count: Int) {
+        promptTokenCount = count
     }
 
-    public var getGeneratedTokenCount: Int {
-        return generatedTokenCount
+    /// Records the number of tokens produced during generation.
+    public func recordGeneratedTokens(_ count: Int) {
+        generatedTokenCount = count
     }
 
     // MARK: - Computed Metrics (from StatsStorage)
@@ -236,7 +237,6 @@ public final class PerformanceMetrics {
         endInstant = nil
         promptTokenCount = 0
         generatedTokenCount = 0
-        totalTokenCount = 0
         // Also reset StatsStorage since this is a full reset
         StatsStorage.shared.reset()
     }

swift/Sources/CoreAILanguageModels/Profiling/Timing.swift

@@ -9,6 +9,10 @@ import Foundation
 
 /// Extension to convert Duration to common time units.
 ///
+/// - Note: This is an internal implementation detail. It is intentionally not
+///   `public`: vending members on a standard-library type we don't own would
+///   pollute `Duration`'s API surface for every client of this library.
+///
 /// Example usage:
 /// ```swift
 /// let start = ContinuousClock.now
@@ -18,13 +22,13 @@ import Foundation
 /// ```
 extension Duration {
     /// Duration in seconds as a Double.
-    public var inSeconds: Double {
+    var inSeconds: Double {
         let (secs, attoseconds) = self.components
         return Double(secs) + Double(attoseconds) / 1e18
     }
 
     /// Duration in milliseconds as a Double.
-    public var inMilliseconds: Double {
+    var inMilliseconds: Double {
         inSeconds * 1000.0
     }
 }

swift/Sources/CoreAIShared/Logger/Logger.swift

@@ -11,7 +11,7 @@ import Synchronization
 public struct CLILogger {
     private static let _level = Atomic<Int>(0)
 
-    static var level: Int {
+    public static var level: Int {
         get {
             _level.load(ordering: .acquiring)
         }
@@ -21,10 +21,6 @@ public struct CLILogger {
         }
     }
 
-    public static func setLevel(to level: Int) {
-        Self.level = level
-    }
-
     /// Performs logging if enabled for the requested level.
     /// - Parameters:
     ///   - message: The message to log.

swift/Sources/Tools/image-segmenter/ImageSegmentationRunnerMain.swift

@@ -157,7 +157,7 @@ struct ImageSegmenterCLI: AsyncParsableCommand {
             throw ValidationError("--image is required.")
         }
         if verbose {
-            CLILogger.setLevel(to: 1)
+            CLILogger.level = 1
         }
 
         if verbose { print("Creating image segmenter...") }

swift/Sources/Tools/llm-runner/LLMRunnerMain.swift

@@ -221,7 +221,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
 
     func run() async throws {
         let verboseLevel = max(self.verboseLevel ?? 0, verbose ? 1 : 0)
-        CLILogger.setLevel(to: verboseLevel)
+        CLILogger.level = verboseLevel
 
         let resolver = ModelPaths()
         let resolvedPath = try validateAndResolveModelPath(resolver: resolver)
@@ -308,7 +308,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
 
         // Set up verbose logging environment variable first
         let verboseLevel = max(self.verboseLevel ?? 0, verbose ? 1 : 0)
-        CLILogger.setLevel(to: verboseLevel)
+        CLILogger.level = verboseLevel
 
         // Bridge hidden CLI overrides to environment variables read by the Core AI engine
         if let b = bucketSize {
@@ -530,7 +530,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
         CLILogger.log("Text generator built successfully", component: "Main")
 
         // Apply chat template and count tokens for metrics
-        await PerformanceMetrics.shared.setPromptTokenCount(promptTokens.count)
+        await PerformanceMetrics.shared.recordPromptTokens(promptTokens.count)
 
         CLILogger.log("Generating text...", component: "Main")
         CLILogger.log("Input: \(displayPrompt)", component: "Main")
@@ -632,7 +632,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
             InstrumentsProfiler.endDecoding(signpostID: decodingID)
 
             // Generated token count is already set by the decoding strategy
-            let generatedTokenCount = await PerformanceMetrics.shared.getGeneratedTokenCount
+            let generatedTokenCount = await PerformanceMetrics.shared.generatedTokenCount
             InstrumentsProfiler.endInference(generatedTokens: generatedTokenCount, signpostID: inferenceID)
 
             // End overall timing now that core inference is complete
@@ -924,7 +924,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
         let inferenceID = InstrumentsProfiler.beginInference(
             promptTokens: vlmTokens.count, maxTokens: maxTokens)
 
-        await PerformanceMetrics.shared.setPromptTokenCount(vlmTokens.count)
+        await PerformanceMetrics.shared.recordPromptTokens(vlmTokens.count)
 
         let tokenStream = try await vlmEngine.generate(
             with: embeddedInput,
@@ -993,7 +993,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
         // Record generation stats
         InstrumentsProfiler.endInference(
             generatedTokens: generatedTokens.count, signpostID: inferenceID)
-        await PerformanceMetrics.shared.setGeneratedTokenCount(generatedTokens.count)
+        await PerformanceMetrics.shared.recordGeneratedTokens(generatedTokens.count)
         await PerformanceMetrics.shared.endOverallTiming()
         await PerformanceMetrics.shared.printSummary(verbose: CLILogger.isVerbose)
 

swift/Tests/LanguageModelsTests/PerformanceMetricsTests.swift

@@ -37,12 +37,12 @@ struct PerformanceMetricsTests {
         metrics.reset()
 
         metrics.startOverallTiming()
-        metrics.setPromptTokenCount(100)
-        metrics.setGeneratedTokenCount(50)
+        metrics.recordPromptTokens(100)
+        metrics.recordGeneratedTokens(50)
         metrics.reset()
 
         #expect(metrics.totalTime == 0)
-        #expect(metrics.getGeneratedTokenCount == 0)
+        #expect(metrics.generatedTokenCount == 0)
         #expect(metrics.modelLoadTime == 0)
     }
 
@@ -51,8 +51,8 @@ struct PerformanceMetricsTests {
         let metrics = PerformanceMetrics.shared
         metrics.reset()
 
-        metrics.setPromptTokenCount(100)
-        metrics.setGeneratedTokenCount(50)
+        metrics.recordPromptTokens(100)
+        metrics.recordGeneratedTokens(50)
 
         #expect(metrics.promptThroughput == 0)
         #expect(metrics.generationThroughput == 0)