module Rager::Utils::Replicate

   # File lib/rager/utils/replicate.rb
14 def self.download_prediction(prediction_url, key = nil, http_adapter: nil)
15   download_url = get_download_url(prediction_url, key, http_adapter: http_adapter)
16   return nil if download_url.nil?
17 
18   Rager::Utils::Http.download_binary(download_url)
19 end

   # File lib/rager/utils/replicate.rb
22 def self.download_prediction_to_file(prediction_url, key = nil, path:, http_adapter: nil)
23   download_url = get_download_url(prediction_url, key, http_adapter: http_adapter)
24   return nil if download_url.nil?
25 
26   Rager::Utils::Http.download_file(download_url, path)
27 end

   # File lib/rager/utils/replicate.rb
30 def self.get_download_url(prediction_url, key = nil, http_adapter: nil)
31   api_key = ENV["REPLICATE_API_KEY"]
32   raise Rager::Errors::CredentialsError.new("Replicate", env_var: ["REPLICATE_API_KEY"]) if api_key.nil?
33 
34   adapter = http_adapter || Rager.config.http_adapter
35   response = adapter.make_request(
36     Rager::Http::Request.new(
37       url: prediction_url,
38       headers: {"Authorization" => "Bearer #{api_key}", "Content-Type" => "application/json"}
39     )
40   )
41 
42   raise Rager::Errors::HttpError.new(adapter, prediction_url, response.status, body: response.body&.to_s) unless response.success?
43 
44   data = JSON.parse(T.cast(T.must(response.body), String))
45   return nil if ["starting", "processing"].include?(data["status"])
46 
47   if data["status"] == "failed"
48     error_msg = data["error"] || "Prediction failed"
49     raise Rager::Errors::HttpError.new(adapter, prediction_url, 422, body: error_msg)
50   end
51 
52   return nil unless data["status"] == "succeeded"
53 
54   output = data["output"]
55   return nil if output.nil?
56 
57   download_url = if key && output.is_a?(Hash) && output[key]
58     output[key]
59   elsif output.is_a?(Hash)
60     output.values.compact.first
61   elsif output.is_a?(Array) && !output.empty?
62     output.first
63   else
64     output.to_s
65   end
66 
67   return nil if download_url.nil? || download_url.empty?
68   download_url
69 rescue JSON::ParserError => e
70   raise Rager::Errors::ParseError.new("Failed to parse prediction response", details: e.message)
71 end

   # File lib/rager/utils/replicate.rb
74 def self.poll_prediction(prediction_url, key: nil, path: nil, max_attempts: 30, sleep_interval: 10, http_adapter: nil)
75   max_attempts.times do
76     result = if path
77       download_prediction_to_file(prediction_url, key, path: path, http_adapter: http_adapter)
78     else
79       download_prediction(prediction_url, key, http_adapter: http_adapter)
80     end
81     return result unless result.nil?
82     Rager::Utils::Http.sleep(sleep_interval)
83   end
84 
85   raise Rager::Errors::TimeoutError.new(
86     "Replicate prediction polling",
87     timeout_seconds: max_attempts * sleep_interval,
88     attempts: max_attempts,
89     details: "Prediction at #{prediction_url} did not complete within #{max_attempts} attempts"
90   )
91 end