rager 0.2.0 → 0.3.0

data/lib/rager/mesh_gen/providers/replicate.rb

@@ -15,12 +15,6 @@ module Rager
           api_key = options.api_key || ENV["REPLICATE_API_KEY"]
           raise Rager::Errors::MissingCredentialsError.new("Replicate", "REPLICATE_API_KEY") if api_key.nil?
 
-          headers = {
-            "Authorization" => "Bearer #{api_key}",
-            "Content-Type" => "application/json",
-            "Prefer" => "wait"
-          }
-
           body = {
             version: options.version,
             input: {
@@ -31,29 +25,29 @@ module Rager
               save_gaussian_ply: true,
               ss_sampling_steps: 38
             }
-          }
-          body[:input][:seed] = options.seed unless options.seed.nil?
+          }.tap do |b|
+            b[:input][:seed] = options.seed if options.seed
+          end
 
           request = Rager::Http::Request.new(
             url: "https://api.replicate.com/v1/predictions",
             verb: Rager::Http::Verb::Post,
-            headers: headers,
+            headers: {
+              "Authorization" => "Bearer #{api_key}",
+              "Content-Type" => "application/json"
+            },
             body: body.to_json
           )
 
           response = Rager.config.http_adapter.make_request(request)
           response_body = T.cast(T.must(response.body), String)
 
-          raise Rager::Errors::HttpError.new(Rager.config.http_adapter, response.status, response_body) unless [200,
-            201, 202].include?(response.status)
+          raise Rager::Errors::HttpError.new(Rager.config.http_adapter, response.status, response_body) unless [200, 201, 202].include?(response.status)
 
-          begin
-            json = JSON.parse(response_body)
-            model_file = json.dig("output", "model_file")
-            model_file || json.fetch("urls").fetch("get")
-          rescue JSON::ParserError, KeyError => e
-            raise Rager::Errors::ParseError.new(e.message, response_body)
-          end
+          json = JSON.parse(response_body)
+          json.fetch("urls").fetch("get")
+        rescue JSON::ParserError, KeyError => e
+          raise Rager::Errors::ParseError.new(e.message, response_body)
         end
       end
     end

data/lib/rager/mesh_gen.rb

@@ -9,7 +9,7 @@ module Rager
 
     sig do
       params(
-        image_url: String,
+        image_url: Rager::Types::MeshGenInput,
         options: Rager::MeshGen::Options
       ).returns(Rager::Types::MeshGenOutput)
     end

data/lib/rager/options.rb

@@ -8,13 +8,15 @@ module Rager
     extend T::Sig
     extend T::Helpers
     requires_ancestor { T::Struct }
-    interface!
 
-    sig { abstract.returns(T::Hash[String, T.untyped]) }
+    sig { returns(T::Hash[String, T.untyped]) }
     def serialize_safe
+      result = T.cast(self, T::Struct).serialize
+      result["api_key"] = "[REDACTED]" if result.key?("api_key")
+      result
     end
 
-    sig { abstract.void }
+    sig { void }
     def validate
     end
   end

data/lib/rager/rerank/options.rb

@@ -14,17 +14,6 @@ module Rager
       const :model, T.nilable(String)
       const :n, T.nilable(Integer)
       const :api_key, T.nilable(String)
-
-      sig { override.returns(T::Hash[String, T.untyped]) }
-      def serialize_safe
-        result = serialize
-        result["api_key"] = "[REDACTED]" if result.key?("api_key")
-        result
-      end
-
-      sig { override.void }
-      def validate
-      end
     end
   end
 end

data/lib/rager/rerank/providers/cohere.rb

@@ -20,50 +20,36 @@ module Rager
           api_key = options.api_key || ENV["COHERE_API_KEY"]
           raise Rager::Errors::MissingCredentialsError.new("Cohere", "COHERE_API_KEY") if api_key.nil?
 
-          url = options.url || ENV["COHERE_URL"] || "https://api.cohere.com/v2/rerank"
-          model = options.model || "rerank-v3.5"
-
-          headers = {
-            "Content-Type" => "application/json"
-          }
-          headers["Authorization"] = "Bearer #{api_key}" if api_key
-
           body = {
-            model: model,
+            model: options.model || "rerank-v3.5",
             query: query.query,
             documents: query.documents
-          }
-          body[:top_n] = options.n if options.n
+          }.tap do |b|
+            b[:top_n] = options.n if options.n
+          end
+
+          headers = {"Content-Type" => "application/json"}
+          headers["Authorization"] = "Bearer #{api_key}" if api_key
 
           request = Rager::Http::Request.new(
             verb: Rager::Http::Verb::Post,
-            url: url,
+            url: options.url || ENV["COHERE_URL"] || "https://api.cohere.com/v2/rerank",
             headers: headers,
             body: body.to_json
           )
 
-          http_adapter = Rager.config.http_adapter
-          response = http_adapter.make_request(request)
+          response = Rager.config.http_adapter.make_request(request)
           response_body = T.cast(T.must(response.body), String)
 
-          if response.status != 200
-            raise Rager::Errors::HttpError.new(
-              http_adapter,
-              response.status,
-              response_body
-            )
-          end
+          raise Rager::Errors::HttpError.new(Rager.config.http_adapter, response.status, response_body) if response.status != 200
 
           parsed_response = JSON.parse(response_body)
-
           results = parsed_response["results"] || []
-          results.map do |result|
-            index = result["index"]
-            score = result["relevance_score"]
 
+          results.map do |result|
             Rager::Rerank::Result.new(
-              score: score,
-              index: index
+              score: result["relevance_score"],
+              index: result["index"]
             )
           end
         end

data/lib/rager/rerank.rb

@@ -9,7 +9,7 @@ module Rager
 
     sig do
       params(
-        query: Rager::Rerank::Query,
+        query: Rager::Types::RerankInput,
         options: Rager::Rerank::Options
       ).returns(Rager::Types::RerankOutput)
     end

data/lib/rager/result.rb

@@ -2,6 +2,7 @@
 # frozen_string_literal: true
 
 require "json"
+require "logger"
 require "securerandom"
 require "sorbet-runtime"
 
@@ -12,6 +13,9 @@ module Rager
     sig { returns(String) }
     attr_reader :id
 
+    sig { returns(String) }
+    attr_reader :context_id
+
     sig { returns(Rager::Operation) }
     attr_reader :operation
 
@@ -21,6 +25,9 @@ module Rager
     sig { returns(Rager::Types::Input) }
     attr_reader :input
 
+    sig { returns(T.nilable(Rager::Types::Output)) }
+    attr_reader :output
+
     sig { returns(Integer) }
     attr_reader :start_time
 
@@ -28,57 +35,62 @@ module Rager
     attr_reader :end_time
 
     sig { returns(T.nilable(String)) }
-    attr_reader :result_id
+    attr_reader :name
 
-    sig { returns(T.nilable(String)) }
-    attr_reader :context_id
+    sig { returns(T.nilable(T::Array[String])) }
+    attr_reader :iids
 
     sig { returns(T.nilable(String)) }
-    attr_reader :hash
-
-    sig { returns(T.nilable(T::Array[String])) }
-    attr_reader :input_ids
+    attr_reader :error
 
     sig do
       params(
         id: String,
+        context_id: String,
         operation: Rager::Operation,
-        options: Rager::Options,
         input: Rager::Types::Input,
-        output: Rager::Types::Output,
+        output: T.nilable(Rager::Types::Output),
+        options: Rager::Options,
         start_time: Integer,
         end_time: Integer,
-        context_id: T.nilable(String),
-        hash: T.nilable(String),
-        input_ids: T.nilable(T::Array[String])
+        name: T.nilable(String),
+        iids: T.nilable(T::Array[String]),
+        error: T.nilable(String)
       ).void
     end
     def initialize(
       id:,
+      context_id:,
       operation:,
-      options:,
       input:,
       output:,
+      options:,
       start_time:,
       end_time:,
-      context_id: nil,
-      hash: nil,
-      input_ids: nil
+      name: nil,
+      iids: nil,
+      error: nil
     )
       @id = id
+      @context_id = context_id
       @operation = operation
-      @options = options
       @input = input
       @output = output
+      @options = options
       @start_time = start_time
       @end_time = end_time
+      @name = T.let(name, T.nilable(String))
+      @iids = T.let(iids, T.nilable(T::Array[String]))
+      @error = T.let(error, T.nilable(String))
+
       @stream = T.let(nil, T.nilable(Rager::Types::Stream))
       @buffer = T.let([], Rager::Types::Buffer)
       @consumed = T.let(false, T::Boolean)
-      @result_id = T.let(SecureRandom.uuid, T.nilable(String))
-      @context_id = T.let(context_id, T.nilable(String))
-      @hash = T.let(hash, T.nilable(String))
-      @input_ids = T.let(input_ids, T.nilable(T::Array[String]))
+    end
+
+    sig { returns(T::Boolean) }
+    def success?
+      @error.nil?
     end
 
     sig { returns(T::Boolean) }
@@ -88,11 +100,9 @@ module Rager
 
     sig { returns(Rager::Types::Output) }
     def out
-      return @output unless stream?
+      return T.must(@output) unless stream?
       return @buffer.each if @consumed
 
-      log
-
       @stream = Enumerator.new do |yielder|
         T.cast(@output, Rager::Types::Stream)
           .each { |message_delta|
@@ -157,49 +167,41 @@ module Rager
       end
     end
 
-    sig { returns(Rager::Types::NonStreamOutput) }
+    sig { returns(T.nilable(Rager::Types::NonStreamOutput)) }
     def serialize_output
-      if @consumed
-        mat
-      elsif stream?
-        "[STREAM]"
-      else
-        T.cast(@output, Rager::Types::NonStreamOutput)
-      end
+      return nil unless success?
+      return @consumed ? mat : "[STREAM]" if stream?
+      T.cast(@output, Rager::Types::NonStreamOutput)
     end
 
     sig { returns(T::Hash[String, T.untyped]) }
     def to_h
       {
         id: @id,
+        context_id: @context_id,
         operation: @operation.serialize,
-        options: @options.serialize_safe,
         input: serialize_input,
-        output: if @consumed
-                  mat
-                elsif stream?
-                  "[STREAM]"
-                else
-                  @output
-                end,
+        output: serialize_output,
+        options: @options.serialize_safe,
         start_time: @start_time,
         end_time: @end_time,
-        result_id: @result_id,
-        context_id: @context_id,
-        hash: @hash,
-        input_ids: @input_ids
+        name: @name,
+        iids: @iids,
+        error: @error
       }
     end
 
     sig { void }
     def log
-      return unless Rager.config.logger
+      return unless Rager.config.logger_type
 
       json = to_h.to_json
 
-      case Rager.config.logger
+      logger = Rager.config.logger
+
+      case Rager.config.logger_type
       when Rager::Logger::Stdout
-        puts "\nLOG: #{json}"
+        success? ? logger.info(json) : logger.error(json)
       when Rager::Logger::Remote
         http_adapter = Rager.config.http_adapter
         url = Rager.config.url
@@ -220,9 +222,7 @@ module Rager
 
           response = http_adapter.make_request(request)
 
-          unless response.status >= 200 && response.status < 300
-            warn "Remote log failed: \\#{response.status} \\#{response.body}"
-          end
+          logger.warn "Remote log failed: \\#{response.status} \\#{response.body}" unless response.success?
         end
       end
     end

data/lib/rager/search/options.rb

@@ -12,17 +12,6 @@ module Rager
       const :provider, String, default: "brave"
       const :n, T.nilable(Integer)
       const :api_key, T.nilable(String)
-
-      sig { override.returns(T::Hash[String, T.untyped]) }
-      def serialize_safe
-        result = serialize
-        result["api_key"] = "[REDACTED]" if result.key?("api_key")
-        result
-      end
-
-      sig { override.void }
-      def validate
-      end
     end
   end
 end

data/lib/rager/search/providers/brave.rb

@@ -21,54 +21,37 @@ module Rager
           api_key = options.api_key || ENV["BRAVE_API_KEY"]
           raise Rager::Errors::MissingCredentialsError.new("Brave", "BRAVE_API_KEY") if api_key.nil?
 
-          params = {"q" => query}
-          params["count"] = options.n.to_s if options.n
-
-          query_string = URI.encode_www_form(params)
-          url = "https://api.search.brave.com/res/v1/web/search?#{query_string}"
-
-          headers = {
-            "Accept" => "application/json",
-            "Accept-Encoding" => "gzip",
-            "x-subscription-token" => api_key
-          }
+          params = {"q" => query}.tap do |p|
+            p["count"] = options.n.to_s if options.n
+          end
 
           request = Rager::Http::Request.new(
             verb: Rager::Http::Verb::Get,
-            url: url,
-            headers: headers
+            url: "https://api.search.brave.com/res/v1/web/search?#{URI.encode_www_form(params)}",
+            headers: {
+              "Accept" => "application/json",
+              "Accept-Encoding" => "gzip",
+              "x-subscription-token" => api_key
+            }
           )
 
-          http_adapter = Rager.config.http_adapter
-          response = http_adapter.make_request(request)
+          response = Rager.config.http_adapter.make_request(request)
           response_body = T.cast(T.must(response.body), String)
 
-          if response.status != 200
-            raise Rager::Errors::HttpError.new(
-              http_adapter,
-              response.status,
-              response_body
-            )
-          end
+          raise Rager::Errors::HttpError.new(Rager.config.http_adapter, response.status, response_body) if response.status != 200
 
           parsed_response = JSON.parse(response_body)
+          web_results = parsed_response.dig("web", "results") || []
 
-          results = []
-
-          if parsed_response["web"] && parsed_response["web"]["results"]
-            web_results = parsed_response["web"]["results"]
-            web_results.each do |result|
-              next unless result["title"] && result["url"] && result["description"]
+          web_results.filter_map do |result|
+            next unless result["title"] && result["url"] && result["description"]
 
-              results << Rager::Search::Result.new(
-                title: result["title"],
-                url: result["url"],
-                description: result["description"]
-              )
-            end
+            Rager::Search::Result.new(
+              title: result["title"],
+              url: result["url"],
+              description: result["description"]
+            )
           end
-
-          results
         end
       end
     end

data/lib/rager/search.rb

@@ -9,7 +9,7 @@ module Rager
 
     sig do
       params(
-        query: String,
+        query: Rager::Types::SearchInput,
         options: Rager::Search::Options
       ).returns(Rager::Types::SearchOutput)
     end

data/lib/rager/template/options.rb

@@ -8,15 +8,6 @@ module Rager
       extend T::Sig
 
       const :provider, String, default: "erb"
-
-      sig { override.returns(T::Hash[String, T.untyped]) }
-      def serialize_safe
-        serialize.transform_keys(&:to_s)
-      end
-
-      sig { override.void }
-      def validate
-      end
     end
   end
 end

data/lib/rager/template.rb

@@ -9,7 +9,7 @@ module Rager
 
     sig do
       params(
-        input: Rager::Template::Input,
+        input: Rager::Types::TemplateInput,
         options: Rager::Template::Options
       ).returns(Rager::Types::TemplateOutput)
     end

data/lib/rager/types.rb

@@ -52,7 +52,6 @@ module Rager
     }
 
     ChatNonStreamOutput = T.type_alias { ChatNonStream }
-
     NonStreamOutput = T.type_alias {
       T.any(
         ChatNonStreamOutput,

data/lib/rager/utils/http.rb

@@ -0,0 +1,39 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+require "uri"
+
+module Rager
+  module Utils
+    module Http
+      extend T::Sig
+
+      sig { params(url: String, path: T.nilable(String), headers: T.nilable(T::Hash[String, String])).returns(String) }
+      def self.download_file(url, path = nil, headers = nil)
+        path ||= begin
+          uri_path = URI.parse(url).path
+          if uri_path.nil? || uri_path.empty?
+            "output"
+          else
+            filename = File.basename(uri_path)
+            filename.empty? ? "output" : filename
+          end
+        end
+
+        response = Rager.config.http_adapter.make_request(
+          Rager::Http::Request.new(url: url, headers: headers || {})
+        )
+
+        raise Rager::Errors::HttpError.new(Rager.config.http_adapter, response.status, T.cast(response.body, T.nilable(String))) unless response.success?
+
+        File.binwrite(path, T.cast(T.must(response.body), String))
+        path
+      rescue Rager::Errors::HttpError
+        raise
+      rescue => e
+        raise Rager::Errors::HttpError.new(Rager.config.http_adapter, 500, "Download failed: #{e.message}")
+      end
+    end
+  end
+end

data/lib/rager/utils/replicate.rb

@@ -0,0 +1,72 @@
+# 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), path: T.nilable(String)).returns(T.nilable(String)) }
+      def self.download_prediction(prediction_url, key = nil, path = nil)
+        api_key = ENV["REPLICATE_API_KEY"]
+        raise Rager::Errors::MissingCredentialsError.new("Replicate", "REPLICATE_API_KEY") if api_key.nil?
+
+        response = Rager.config.http_adapter.make_request(
+          Rager::Http::Request.new(
+            url: prediction_url,
+            headers: {"Authorization" => "Bearer #{api_key}", "Content-Type" => "application/json"}
+          )
+        )
+
+        raise Rager::Errors::HttpError.new(Rager.config.http_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(Rager.config.http_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?
+
+        Rager::Utils::Http.download_file(download_url, path)
+      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).returns(T.nilable(String)) }
+      def self.poll_prediction(prediction_url, key: nil, path: nil, max_attempts: 30, sleep_interval: 10)
+        max_attempts.times do
+          result = download_prediction(prediction_url, key, path)
+          return result unless result.nil?
+          sleep(sleep_interval)
+        end
+
+        raise Rager::Errors::HttpError.new(Rager.config.http_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.0"
+  VERSION = "0.3.0"
 end