Add progress property to Prediction

replicate · Feb 19, 2024 · 260b2f8 · 260b2f8
1 parent 1893a95
commit 260b2f8
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 1 deletion.
diff --git a/Sources/Replicate/Prediction.swift b/Sources/Replicate/Prediction.swift
@@ -1,8 +1,11 @@
 import struct Foundation.Date
 import struct Foundation.TimeInterval
 import struct Foundation.URL
+import class Foundation.Progress
 import struct Dispatch.DispatchTime
 
+import RegexBuilder
+
 /// A prediction with unspecified inputs and outputs.
 public typealias AnyPrediction = Prediction<[String: Value], Value>
 
@@ -81,6 +84,32 @@ public struct Prediction<Input, Output>: Identifiable where Input: Codable, Outp
 
     // MARK: -
 
+    @available(macOS 13.0, iOS 16.0, watchOS 9.0, tvOS 16.0, *)
+    public var progress: Progress? {
+        guard let logs = self.logs else { return nil }
+
+        let regex: Regex = #/^\s*(\d+)%\s*\|.+?\|\s*(\d+)\/(\d+)/#
+
+        let lines = logs.split(separator: "\n")
+        guard !lines.isEmpty else { return nil }
+
+        for line in lines.reversed() {
+            let lineString = String(line).trimmingCharacters(in: .whitespaces)
+            if let match = try? regex.firstMatch(in: lineString),
+               let current = Int64(match.output.2),
+               let total = Int64(match.output.3)
+            {
+                let progress = Progress(totalUnitCount: total)
+                progress.completedUnitCount = current
+                return progress
+            }
+        }
+
+        return nil
+    }
+
+    // MARK: -
+
     /// Wait for the prediction to complete.
     ///
     /// - Parameters:

diff --git a/Tests/ReplicateTests/ClientTests.swift b/Tests/ReplicateTests/ClientTests.swift
@@ -82,6 +82,8 @@ final class ClientTests: XCTestCase {
         XCTAssertEqual(prediction.status, .succeeded)
         XCTAssertEqual(prediction.createdAt.timeIntervalSinceReferenceDate, 672703986.224, accuracy: 1)
         XCTAssertEqual(prediction.urls["cancel"]?.absoluteString, "https://api.replicate.com/v1/predictions/ufawqhfynnddngldkgtslldrkq/cancel")
+        XCTAssertEqual(prediction.progress?.completedUnitCount, 5)
+        XCTAssertEqual(prediction.progress?.totalUnitCount, 5)
     }
 
     func testCancelPrediction() async throws {

diff --git a/Tests/ReplicateTests/Helpers/MockURLProtocol.swift b/Tests/ReplicateTests/Helpers/MockURLProtocol.swift
@@ -142,7 +142,7 @@ class MockURLProtocol: URLProtocol {
                       },
                       "output": ["Hello, Alice!"],
                       "error": null,
-                      "logs": "",
+                      "logs": "Using seed: 12345,\n0%|          | 0/5 [00:00<?, ?it/s]\n20%|██        | 1/5 [00:00<00:01, 21.38it/s]\n40%|████▍     | 2/5 [00:01<00:01, 22.46it/s]\n60%|████▍     | 3/5 [00:01<00:01, 22.46it/s]\n 80%|████████  | 4/5 [00:01<00:00, 22.86it/s]\n100%|██████████| 5/5 [00:02<00:00, 22.26it/s]",
                       "metrics": {
                         "predict_time": 10.0
                       }