Add automatic fallback to polling API for replicate requests (#104)

Author: lzell

README.md

@@ -13,7 +13,7 @@ let aiproxyLogger = Logger(
 public struct AIProxy {
 
     /// The current sdk version
-    public static let sdkVersion = "0.69.0"
+    public static let sdkVersion = "0.70.0"
 
     /// - Parameters:
     ///   - partialKey: Your partial key is displayed in the AIProxy dashboard when you submit your provider's key.
@@ -379,13 +379,15 @@ public struct AIProxy {
     /// - Parameters:
     ///   - unprotectedAPIKey: Your Replicate API key
     /// - Returns: An instance of ReplicateService configured and ready to make requests
+    #if false
     public static func replicateDirectService(
         unprotectedAPIKey: String
     ) -> ReplicateService {
         return ReplicateDirectService(
             unprotectedAPIKey: unprotectedAPIKey
         )
     }
+    #endif
 
     /// AIProxy's ElevenLabs service
     ///

Sources/AIProxy/AIProxy.swift

@@ -21,7 +21,7 @@ public enum ReplicateError: LocalizedError {
         case .missingModelURL:
             return "The replicate model does not contain a URL"
         case .reachedRetryLimit:
-            return "Reached replicate polling retry limit"
+            return "Reached secondsToWait without the prediction completing"
         }
     }
 }

Sources/AIProxy/Replicate/ReplicateDirectService.swift

@@ -7,17 +7,23 @@
 
 import Foundation
 
-struct ReplicateFileUploadRequestBody: MultipartFormEncodable {
+internal struct ReplicateFileUploadRequestBody: MultipartFormEncodable {
 
-    let fileData: Data
+    /// The binary contents of the file
+    let contents: Data
+
+    /// The file mime type
+    let contentType: String
+
+    /// The name of the file. I believe this does not get preserved on replicate's CDN. Can it be removed?
     let fileName: String
 
     var formFields: [FormField] {
         return [
             .fileField(
                 name: "content",
-                content: self.fileData,
-                contentType: "application/zip",
+                content: self.contents,
+                contentType: self.contentType,
                 filename: self.fileName
             )
         ]

Sources/AIProxy/Replicate/ReplicateError.swift

@@ -7,13 +7,21 @@
 
 import Foundation
 
-/// The request body for creating a Replicate prediction:
-/// https://replicate.com/docs/reference/http#create-a-prediction
+/// The request body for creating a Replicate prediction.
+///
+/// This type is used for both community models and official models.
+/// When using with an official model, the `version` property can remain `nil`.
+///
+/// Community model reference: https://replicate.com/docs/reference/http#predictions.create
+/// Official model reference: https://replicate.com/docs/reference/http#models.predictions.create
 public struct ReplicatePredictionRequestBody: Encodable {
+
     /// The replicate input schema, for example ReplicateSDXLInputSchema
+    /// TThe input schema depends on what model you are running. To see the available inputs, click the "API" tab on the model you are running or get the model version and look at its `openapi_schema` property. For example, `stability-ai/sdxl` takes `prompt` as an input.
     public let input: Encodable
 
-    /// The version of the model to run
+    /// You do not need to set this field if you are using an official model.
+    /// For community models, set it to the ID of the model version that you want to run.
     public let version: String?
 
     public init(

Sources/AIProxy/Replicate/ReplicateFileUploadRequestBody.swift

@@ -1,34 +1,62 @@
 //
-//  ReplicatePredictionResponseBody.swift
+//  ReplicatePrediction.swift
 //
 //
 //  Created by Lou Zell on 8/25/24.
 //
 
 import Foundation
 
-/// Response body for a Replicate prediction.
-/// This format is used for both the "create a predition" and "get a prediction" endpoints:
+public typealias ReplicatePredictionResponseBody = ReplicatePrediction
+
+/// Represents the current state of a replicate prediction
+///
+/// This type is used for both the "create a predition" and "get a prediction" endpoints:
 ///     https://replicate.com/docs/reference/http#get-a-prediction
 ///     https://replicate.com/docs/reference/http#create-a-prediction
-public struct ReplicatePredictionResponseBody<T: Decodable>: Decodable {
+///
+/// And it is used for both the sync and polling API:
+///     https://replicate.com/docs/topics/predictions/create-a-prediction#sync-mode
+///     https://replicate.com/docs/topics/predictions/create-a-prediction#polling
+public struct ReplicatePrediction<Output: Decodable>: Decodable {
 
     /// ISO8601 date stamp of when the prediction completed
     public let completedAt: String?
 
+    /// ISO8601 date stamp of when the prediction was created
+    public let createdAt: String?
+
+    /// https://replicate.com/docs/topics/predictions/data-retention
+    public let dataRemoved: Bool?
+
     /// In the case of failure, error will contain the error encountered during the prediction
     public let error: String?
 
+    /// ID of the prediction
+    public let id: String?
+
+    /// Prediction logs
+    public let logs: String?
+
+    /// How long the prediction took
+    public let metrics: Metrics?
+
+    /// The model being run
+    public let model: String?
+
     /// The output adheres to Replicate's "output schema" structure.
     /// Schemas can be found at the Replicate model's detail page by tapping on `API > Schema > Output Schema`.
     /// In the case of SDXL, the output is an array of URLs, which you can see here:
     /// https://replicate.com/stability-ai/sdxl/api/schema#output-schema
-    public let output: T?
+    public let output: Output?
 
     /// ISO8601 timestamp of start of prediction
     public let startedAt: String?
 
-    /// One of `starting`, `processing`, `succeeded`, `failed`, `canceled`
+    /// One of `starting`, `processing`, `succeeded`, `failed`, `canceled`.
+    ///
+    /// In the `succeeded` case, the `output` property on this type will be an object containing the output of the model.
+    /// In the `failed` case, `error` property on this type will contain the error encountered during the prediction.
     public let status: Status?
 
     /// URLs to cancel the prediction or get the result from the prediction
@@ -37,9 +65,20 @@ public struct ReplicatePredictionResponseBody<T: Decodable>: Decodable {
     /// The version of the model that ran
     public let version: String?
 
+    /// For compatibility with an older release of the libary
+    public var predictionResultURL: URL? {
+        return urls?.get
+    }
+
     private enum CodingKeys: String, CodingKey {
         case completedAt = "completed_at"
+        case createdAt = "created_at"
+        case dataRemoved = "data_removed"
         case error
+        case id
+        case logs
+        case metrics
+        case model
         case output
         case startedAt = "started_at"
         case status
@@ -48,19 +87,42 @@ public struct ReplicatePredictionResponseBody<T: Decodable>: Decodable {
     }
 }
 
-extension ReplicatePredictionResponseBody {
+extension ReplicatePrediction {
     public struct ActionURLs: Decodable {
         public let cancel: URL?
         public let get: URL?
     }
 }
 
-extension ReplicatePredictionResponseBody {
+extension ReplicatePrediction {
+    public struct Metrics: Decodable {
+        public let predictTime: Double
+
+        enum CodingKeys: String, CodingKey {
+            case predictTime = "predict_time"
+        }
+    }
+}
+
+extension ReplicatePrediction {
     public enum Status: String, Decodable {
+        /// The prediction is starting up. If this status lasts longer than a few seconds, then it's typically because a new worker is being started to run the prediction.
         case starting
+
+        /// The `predict()` method of the model is currently running.
         case processing
+
+        /// The prediction completed successfully.
         case succeeded
+
+        /// The prediction encountered an error during processing.
         case failed
+
+        /// The prediction was canceled by its creator
         case canceled
+
+        var isTerminal: Bool {
+            return [.succeeded, .failed, .canceled].contains(self)
+        }
     }
 }