rager 0.1.0

data/lib/rager/errors/missing_credentials_error.rb

@@ -0,0 +1,19 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  module Errors
+    class MissingCredentialsError < Rager::Error
+      extend T::Sig
+
+      sig { params(provider_name: String, env_var: T.nilable(String)).void }
+      def initialize(provider_name, env_var = nil)
+        message = "Missing credentials for provider #{provider_name}"
+        message += " -- attempted lookup with #{env_var}" if env_var
+        super(message)
+      end
+    end
+  end
+end

data/lib/rager/errors/options_error.rb

@@ -0,0 +1,19 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  module Errors
+    class OptionsError < Rager::Error
+      extend T::Sig
+
+      sig { params(invalid_keys: T::Array[String], description: T.nilable(String)).void }
+      def initialize(invalid_keys:, description: nil)
+        message = "Invalid keys #{invalid_keys.join(", ")}"
+        message += " -- #{description})" if description
+        super(message)
+      end
+    end
+  end
+end

data/lib/rager/errors/parse_error.rb

@@ -0,0 +1,19 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  module Errors
+    class ParseError < Rager::Error
+      extend T::Sig
+
+      sig { params(description: String, body: T.nilable(String)).void }
+      def initialize(description, body = nil)
+        message = description
+        message += " -- #{body}" if body
+        super(message)
+      end
+    end
+  end
+end

data/lib/rager/errors/unknown_provider_error.rb

@@ -0,0 +1,17 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  module Errors
+    class UnknownProviderError < Rager::Error
+      extend T::Sig
+
+      sig { params(operation: Rager::Operation, key: String).void }
+      def initialize(operation, key)
+        super("Unknown provider #{key} for operation #{operation.serialize}")
+      end
+    end
+  end
+end

data/lib/rager/http/adapters/abstract.rb

@@ -0,0 +1,20 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  module Http
+    module Adapters
+      class Abstract
+        extend T::Sig
+        extend T::Helpers
+        abstract!
+
+        sig { abstract.params(request: Rager::Http::Request).returns(Rager::Http::Response) }
+        def make_request(request)
+        end
+      end
+    end
+  end
+end

data/lib/rager/http/adapters/async_http.rb

@@ -0,0 +1,65 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  module Http
+    module Adapters
+      class AsyncHttp < Rager::Http::Adapters::Abstract
+        extend T::Sig
+
+        sig { void }
+        def initialize
+          require "async/http"
+
+          @internet = T.let(Async::HTTP::Internet.new, Async::HTTP::Internet)
+        end
+
+        sig {
+          override.params(
+            request: Rager::Http::Request
+          ).returns(Rager::Http::Response)
+        }
+        def make_request(request)
+          response = @internet.call(
+            request.verb.serialize,
+            request.url,
+            request.headers.to_a,
+            request.body
+          )
+
+          body = if response.body.nil?
+            nil
+          elsif (response.headers["Transfer-Encoding"]&.downcase == "chunked") ||
+              response.headers["content-type"]&.downcase&.include?("text/event-stream")
+            body_enum(response)
+          else
+            response.body.read
+          end
+
+          Response.new(
+            status: response.status,
+            headers: response.headers.to_h,
+            body: body
+          )
+        end
+
+        private
+
+        sig {
+          params(
+            response: Async::HTTP::Protocol::Response
+          ).returns(T::Enumerator[String])
+        }
+        def body_enum(response)
+          Enumerator.new do |yielder|
+            response.body.each { |chunk| yielder << chunk.force_encoding("UTF-8") }
+          ensure
+            response.close
+          end
+        end
+      end
+    end
+  end
+end

data/lib/rager/http/adapters/mock.rb

@@ -0,0 +1,138 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "fileutils"
+require "json"
+require "sorbet-runtime"
+
+module Rager
+  module Http
+    module Adapters
+      class Mock < Rager::Http::Adapters::Abstract
+        extend T::Sig
+
+        Cache = T.type_alias { T::Hash[String, T::Hash[String, T.untyped]] }
+
+        sig do
+          params(
+            test_file_path: String,
+            fallback_adapter: T.nilable(Rager::Http::Adapters::Abstract),
+            chunk_delimiter: T.nilable(String)
+          ).void
+        end
+        def initialize(
+          test_file_path,
+          fallback_adapter = nil,
+          chunk_delimiter = nil
+        )
+          @test_file_path = T.let(test_file_path, String)
+          @fallback_adapter = T.let(fallback_adapter || Rager::Http::Adapters::AsyncHttp.new, Rager::Http::Adapters::Abstract)
+          @cache = T.let(load_cache, Cache)
+        end
+
+        sig { override.params(request: Rager::Http::Request).returns(Rager::Http::Response) }
+        def make_request(request)
+          key = request.serialize.to_json
+          cached_entry = @cache[key]
+          if cached_entry
+            build_response_from_cache(cached_entry)
+          else
+            fetch_and_cache_response(request, key)
+          end
+        end
+
+        sig {
+          params(
+            request: Rager::Http::Request,
+            key: String
+          ).returns(Rager::Http::Response)
+        }
+        def fetch_and_cache_response(request, key)
+          response = @fallback_adapter.make_request(request)
+
+          serialized_response = T.let({
+            "status" => response.status,
+            "headers" => response.headers
+          }, T::Hash[String, T.untyped])
+
+          if response.body.is_a?(Enumerator)
+            chunks = T.let([], T::Array[String])
+
+            T.cast(response.body, T::Enumerator[String]).each do |chunk|
+              chunks << chunk
+            end
+
+            serialized_response["body"] = chunks
+            serialized_response["is_stream"] = true
+
+            response_body = Enumerator.new do |yielder|
+              chunks.each { |chunk| yielder << chunk }
+            end
+          else
+            serialized_response["body"] = response.body
+            serialized_response["is_stream"] = false
+            response_body = response.body
+          end
+
+          @cache[key] = serialized_response
+
+          save_cache
+
+          Rager::Http::Response.new(
+            status: response.status,
+            headers: response.headers,
+            body: response_body
+          )
+        end
+
+        sig {
+          params(
+            entry: T::Hash[String, T.untyped]
+          ).returns(Rager::Http::Response)
+        }
+        def build_response_from_cache(entry)
+          body = entry["body"]
+          is_stream = entry["is_stream"] || false
+
+          body = if is_stream
+            T.cast(body, T::Array[String]).to_enum
+          else
+            body
+          end
+
+          Rager::Http::Response.new(
+            status: T.cast(entry["status"], Integer),
+            headers: T.cast(entry["headers"], T::Hash[String, String]),
+            body: body
+          )
+        end
+
+        sig { void }
+        def create_file_if_not_exists
+          FileUtils.mkdir_p(File.dirname(@test_file_path))
+          File.write(@test_file_path, "{}") unless File.exist?(@test_file_path)
+        end
+
+        sig { returns(Cache) }
+        def load_cache
+          create_file_if_not_exists
+          JSON.parse(File.read(@test_file_path, encoding: "UTF-8"))
+        end
+
+        sig { void }
+        def save_cache
+          current_file_content =
+            if File.exist?(@test_file_path)
+              JSON.parse(File.read(@test_file_path, encoding: "UTF-8"))
+            else
+              {}
+            end
+
+          output = current_file_content.merge(@cache)
+          json_string = JSON.generate(output).force_encoding("UTF-8")
+          File.write(@test_file_path, json_string, encoding: "UTF-8")
+        end
+      end
+    end
+  end
+end

data/lib/rager/http/request.rb

@@ -0,0 +1,15 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  module Http
+    class Request < T::Struct
+      const :url, String
+      const :verb, Rager::Http::Verb, default: Rager::Http::Verb::Get
+      const :headers, T::Hash[String, String]
+      const :body, T.nilable(String)
+    end
+  end
+end

data/lib/rager/http/response.rb

@@ -0,0 +1,14 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  module Http
+    class Response < T::Struct
+      const :status, Integer
+      const :headers, T::Hash[String, T.any(String, T::Array[String])]
+      const :body, T.nilable(T.any(String, T::Enumerator[String]))
+    end
+  end
+end

data/lib/rager/http/verb.rb

@@ -0,0 +1,20 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  module Http
+    class Verb < T::Enum
+      enums do
+        Get = new("GET")
+        Post = new("POST")
+        Put = new("PUT")
+        Patch = new("PATCH")
+        Delete = new("DELETE")
+        Head = new("HEAD")
+        Options = new("OPTIONS")
+      end
+    end
+  end
+end

data/lib/rager/image_gen/options.rb

@@ -0,0 +1,29 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  module ImageGen
+    class Options < T::Struct
+      extend T::Sig
+      include Rager::Options
+
+      const :provider, String, default: "replicate"
+      const :model, String, default: "black-forest-labs/flux-schnell"
+      const :api_key, T.nilable(String)
+      const :seed, T.nilable(Integer)
+
+      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/image_gen/providers/abstract.rb

@@ -0,0 +1,25 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  module ImageGen
+    module Providers
+      class Abstract
+        extend T::Sig
+        extend T::Helpers
+        abstract!
+
+        sig do
+          abstract.params(
+            prompt: String,
+            options: Rager::ImageGen::Options
+          ).returns(Rager::Types::ImageGenOutput)
+        end
+        def image_gen(prompt, options)
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,55 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+require "json"
+
+module Rager
+  module ImageGen
+    module Providers
+      class Replicate < Rager::ImageGen::Providers::Abstract
+        extend T::Sig
+
+        sig { override.params(prompt: String, options: Rager::ImageGen::Options).returns(Rager::Types::ImageGenOutput) }
+        def image_gen(prompt, options)
+          url = "https://api.replicate.com/v1/models/#{options.model}/predictions"
+          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 = {
+            input: {
+              prompt: prompt
+            }
+          }
+          body[:input][:seed] = options.seed unless options.seed.nil?
+
+          request = Rager::Http::Request.new(
+            url: url,
+            verb: Rager::Http::Verb::Post,
+            headers: headers,
+            body: body.to_json
+          )
+
+          http_adapter = Rager.config.http_adapter
+          response = http_adapter.make_request(request)
+          response_body = T.cast(T.must(response.body), String)
+
+          raise Rager::Errors::HttpError.new(http_adapter, response.status, response_body) unless [200, 201].include?(response.status)
+
+          begin
+            parsed = JSON.parse(response_body)
+            parsed.fetch("output").first
+          rescue JSON::ParserError, KeyError => e
+            raise Rager::Errors::ParseError.new(e.message, response_body)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/rager/image_gen.rb

@@ -0,0 +1,31 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  module ImageGen
+    extend T::Sig
+
+    sig do
+      params(
+        prompt: String,
+        options: Rager::ImageGen::Options
+      ).returns(Rager::Types::ImageGenOutput)
+    end
+    def self.image_gen(prompt, options = Rager::ImageGen::Options.new)
+      provider = get_provider(options.provider)
+      provider.image_gen(prompt, options)
+    end
+
+    sig { params(key: String).returns(Rager::ImageGen::Providers::Abstract) }
+    def self.get_provider(key)
+      case key.downcase
+      when "replicate"
+        Rager::ImageGen::Providers::Replicate.new
+      else
+        raise Rager::Errors::UnknownProviderError.new(Rager::Operation::ImageGen, key)
+      end
+    end
+  end
+end

data/lib/rager/logger.rb

@@ -0,0 +1,13 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  class Logger < T::Enum
+    enums do
+      Stdout = new("stdout")
+      Remote = new("remote")
+    end
+  end
+end

data/lib/rager/mesh_gen/options.rb

@@ -0,0 +1,30 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  module MeshGen
+    class Options < T::Struct
+      extend T::Sig
+      include Rager::Options
+
+      const :provider, String, default: "replicate"
+      const :model, String, default: "firoz/trellis"
+      const :version, String, default: "4876f2a8da1c544772dffa32e8889da4a1bab3a1f5c1937bfcfccb99ae347251"
+      const :api_key, T.nilable(String)
+      const :seed, T.nilable(Integer)
+
+      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/mesh_gen/providers/abstract.rb

@@ -0,0 +1,25 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  module MeshGen
+    module Providers
+      class Abstract
+        extend T::Sig
+        extend T::Helpers
+        abstract!
+
+        sig do
+          abstract.params(
+            image_url: String,
+            options: Rager::MeshGen::Options
+          ).returns(Rager::Types::MeshGenOutput)
+        end
+        def mesh_gen(image_url, options)
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,61 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+require "json"
+
+module Rager
+  module MeshGen
+    module Providers
+      class Replicate < Rager::MeshGen::Providers::Abstract
+        extend T::Sig
+
+        sig { override.params(image_url: String, options: Rager::MeshGen::Options).returns(Rager::Types::MeshGenOutput) }
+        def mesh_gen(image_url, options)
+          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: {
+              images: [image_url],
+              texture_size: 2048,
+              mesh_simplify: 0.9,
+              generate_model: true,
+              save_gaussian_ply: true,
+              ss_sampling_steps: 38
+            }
+          }
+          body[:input][:seed] = options.seed unless options.seed.nil?
+
+          request = Rager::Http::Request.new(
+            url: "https://api.replicate.com/v1/predictions",
+            verb: Rager::Http::Verb::Post,
+            headers: headers,
+            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)
+
+          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
+        end
+      end
+    end
+  end
+end

data/lib/rager/mesh_gen.rb

@@ -0,0 +1,31 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  module MeshGen
+    extend T::Sig
+
+    sig do
+      params(
+        image_url: String,
+        options: Rager::MeshGen::Options
+      ).returns(Rager::Types::MeshGenOutput)
+    end
+    def self.mesh_gen(image_url, options = Rager::MeshGen::Options.new)
+      provider = get_provider(options.provider)
+      provider.mesh_gen(image_url, options)
+    end
+
+    sig { params(key: String).returns(Rager::MeshGen::Providers::Abstract) }
+    def self.get_provider(key)
+      case key.downcase
+      when "replicate"
+        Rager::MeshGen::Providers::Replicate.new
+      else
+        raise Rager::Errors::UnknownProviderError.new(Rager::Operation::MeshGen, key)
+      end
+    end
+  end
+end

data/lib/rager/operation.rb

@@ -0,0 +1,14 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  class Operation < T::Enum
+    enums do
+      Chat = new("chat")
+      ImageGen = new("image_gen")
+      MeshGen = new("mesh_gen")
+    end
+  end
+end