rager 0.2.1 → 0.4.0

data/lib/rager/utils/replicate.rb

@@ -0,0 +1,97 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+require "json"
+
+module Rager
+  module Utils
+    module Replicate
+      extend T::Sig
+
+      sig { params(prediction_url: String, key: T.nilable(String), http_adapter: T.nilable(Rager::Http::Adapters::Abstract)).returns(T.nilable(String)) }
+      def self.download_prediction(prediction_url, key = nil, http_adapter: nil)
+        download_url = get_download_url(prediction_url, key, http_adapter: http_adapter)
+        return nil if download_url.nil?
+
+        Rager::Utils::Http.download_binary(download_url)
+      rescue
+        nil
+      end
+
+      sig { params(prediction_url: String, key: T.nilable(String), path: String, http_adapter: T.nilable(Rager::Http::Adapters::Abstract)).returns(T.nilable(String)) }
+      def self.download_prediction_to_file(prediction_url, key = nil, path:, http_adapter: nil)
+        download_url = get_download_url(prediction_url, key, http_adapter: http_adapter)
+        return nil if download_url.nil?
+
+        Rager::Utils::Http.download_file(download_url, path)
+      rescue
+        nil
+      end
+
+      sig { params(prediction_url: String, key: T.nilable(String), http_adapter: T.nilable(Rager::Http::Adapters::Abstract)).returns(T.nilable(String)) }
+      def self.get_download_url(prediction_url, key = nil, http_adapter: nil)
+        api_key = ENV["REPLICATE_API_KEY"]
+        raise Rager::Errors::MissingCredentialsError.new("Replicate", "REPLICATE_API_KEY") if api_key.nil?
+
+        adapter = http_adapter || Rager.config.http_adapter
+        response = adapter.make_request(
+          Rager::Http::Request.new(
+            url: prediction_url,
+            headers: {"Authorization" => "Bearer #{api_key}", "Content-Type" => "application/json"}
+          )
+        )
+
+        raise Rager::Errors::HttpError.new(adapter, response.status, T.cast(response.body, T.nilable(String))) unless response.success?
+
+        data = JSON.parse(T.cast(T.must(response.body), String))
+        return nil if ["starting", "processing"].include?(data["status"])
+
+        if data["status"] == "failed"
+          error_msg = data["error"] || "Prediction failed"
+          raise Rager::Errors::HttpError.new(adapter, 422, "Prediction failed: #{error_msg}")
+        end
+
+        return nil unless data["status"] == "succeeded"
+
+        output = data["output"]
+        return nil if output.nil?
+
+        download_url = if key && output.is_a?(Hash) && output[key]
+          output[key]
+        elsif output.is_a?(Hash)
+          output.values.compact.first
+        elsif output.is_a?(Array) && !output.empty?
+          output.first
+        else
+          output.to_s
+        end
+
+        return nil if download_url.nil? || download_url.empty?
+        download_url
+      rescue JSON::ParserError => e
+        raise Rager::Errors::ParseError.new("Failed to parse prediction response", e.message)
+      rescue Rager::Errors::HttpError, Rager::Errors::ParseError, Rager::Errors::MissingCredentialsError
+        raise
+      rescue
+        nil
+      end
+
+      sig { params(prediction_url: String, key: T.nilable(String), path: T.nilable(String), max_attempts: Integer, sleep_interval: Integer, http_adapter: T.nilable(Rager::Http::Adapters::Abstract)).returns(T.nilable(String)) }
+      def self.poll_prediction(prediction_url, key: nil, path: nil, max_attempts: 30, sleep_interval: 10, http_adapter: nil)
+        max_attempts.times do
+          result = if path
+            download_prediction_to_file(prediction_url, key, path: path, http_adapter: http_adapter)
+          else
+            download_prediction(prediction_url, key, http_adapter: http_adapter)
+          end
+          return result unless result.nil?
+          sleep(sleep_interval)
+        end
+
+        adapter = http_adapter || Rager.config.http_adapter
+        raise Rager::Errors::HttpError.new(adapter, 408, "Prediction polling timed out after #{max_attempts} attempts")
+      end
+    end
+  end
+end

data/lib/rager/version.rb

@@ -2,5 +2,5 @@
 # frozen_string_literal: true
 
 module Rager
-  VERSION = "0.2.1"
+  VERSION = "0.4.0"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: rager
 version: !ruby/object:Gem::Version
-  version: 0.2.1
+  version: 0.4.0
 platform: ruby
 authors:
 - mvkvc
@@ -106,6 +106,7 @@ files:
 - lib/rager/http/verb.rb
 - lib/rager/image_gen.rb
 - lib/rager/image_gen/options.rb
+- lib/rager/image_gen/output_format.rb
 - lib/rager/image_gen/providers/abstract.rb
 - lib/rager/image_gen/providers/replicate.rb
 - lib/rager/logger.rb
@@ -133,8 +134,10 @@ files:
 - lib/rager/template/providers/abstract.rb
 - lib/rager/template/providers/erb.rb
 - lib/rager/types.rb
+- lib/rager/utils/http.rb
+- lib/rager/utils/replicate.rb
 - lib/rager/version.rb
-homepage: https://github.com/mvkvc/rager_ruby
+homepage: https://github.com/mvkvc/rager_rb
 licenses:
 - MIT
 metadata: {}
@@ -145,7 +148,7 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: '3.1'
+      version: '3.2'
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="