reranker-ruby 0.1.1 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: bcd9ac5bc1d098f8780a6523e648505dbf0259f4335b07e909898203b9723866
-  data.tar.gz: 12297c4398121fe27b4f0ed148268197f44258dde3b9e5f7efcf305202a5e861
+  metadata.gz: c7fd1607f0f0124cf4a8e2fd10e6a643b738d1c65e555fe73becee2ededd360f
+  data.tar.gz: 354c66d0a2767171e7c4da774a91ee77287ac7b8c97dc02de290f0a1b1b95f19
 SHA512:
-  metadata.gz: 75a6a774a12b00ed332f5f401451a53164f8ef3d8fe23428d70cfc224fa951bf0fd7a655f72178a98194e0578698e9f1de826afcd09b59b2dee2ec2dd0ea494f
-  data.tar.gz: c46089ed7e3ebedc6aea3bd2455bac70cd3b1666b77519a970510a4d4790d2bb3857a5d295d8c35e88358234d39bee5fee8e276fa5b34d8511b47b47a10f87fb
+  metadata.gz: 3317a9db00d34aff5c7d06b73a7d0d2ec67bfd6f180212905d64b0a8d4e3a2426d3cb52564350c91b0741cfcc5f7512d127f021fee43a575073c8e1af9fef610
+  data.tar.gz: eaf726ae758dca2c28ececff924506470a760232396e20536c00f3051da1cdd681433805d60a90c1b8bc7cf5b066fa613e382508ba5e421024876650fd624477

data/lib/reranker_ruby/base.rb

@@ -8,11 +8,18 @@ require "digest"
 module RerankerRuby
   class Error < StandardError; end
   class APIError < Error; end
+  class TimeoutError < Error; end
+  class RateLimitError < APIError; end
+  class ValidationError < Error; end
 
   class Base
-    def initialize(api_key: nil, cache: nil)
+    def initialize(api_key: nil, cache: nil, timeout: nil, max_retries: nil)
       @api_key = api_key
       @cache = cache
+      @timeout = timeout || 30
+      @max_retries = max_retries || 3
+      @connections = {}
+      @conn_mutex = Mutex.new
     end
 
     def rerank(query, documents, top_k: 10)
@@ -22,9 +29,9 @@ module RerankerRuby
     private
 
     def validate_inputs!(query, documents, top_k)
-      raise ArgumentError, "query cannot be nil or empty" if query.nil? || query.to_s.strip.empty?
-      raise ArgumentError, "documents cannot be nil or empty" if documents.nil? || documents.empty?
-      raise ArgumentError, "top_k must be positive" if top_k && top_k <= 0
+      raise ValidationError, "query cannot be nil or empty" if query.nil? || query.to_s.strip.empty?
+      raise ValidationError, "documents cannot be nil or empty" if documents.nil? || documents.empty?
+      raise ValidationError, "top_k must be positive" if top_k && top_k <= 0
     end
 
     def instrument(query:, document_count:, top_k:, &block)
@@ -63,18 +70,32 @@ module RerankerRuby
       result
     end
 
+    def get_connection(uri)
+      host_key = "#{uri.host}:#{uri.port}"
+      @conn_mutex.synchronize do
+        conn = @connections[host_key]
+        if conn && conn.started?
+          return conn
+        end
+
+        http = Net::HTTP.new(uri.host, uri.port)
+        http.use_ssl = uri.scheme == "https"
+        http.open_timeout = @timeout
+        http.read_timeout = @timeout
+        http.write_timeout = @timeout
+        http.keep_alive_timeout = 30
+        http.start
+        @connections[host_key] = http
+        http
+      end
+    end
+
     def post(url, body, headers: {})
       uri = URI.parse(url)
       retries = 0
-      max_retries = 3
 
       begin
-        http = Net::HTTP.new(uri.host, uri.port)
-        http.use_ssl = uri.scheme == "https"
-        http.open_timeout = 30
-        http.read_timeout = 30
-        http.write_timeout = 30
-
+        http = get_connection(uri)
         request = Net::HTTP::Post.new(uri.path)
         request["Content-Type"] = "application/json"
         request["User-Agent"] = "RerankerRuby/#{RerankerRuby::VERSION}"
@@ -83,7 +104,12 @@ module RerankerRuby
 
         response = http.request(request)
 
-        if response.code.to_i == 429 || response.code.to_i >= 500
+        if response.code.to_i == 429
+          retry_after = response["Retry-After"]&.to_f || (2 ** retries)
+          raise RateLimitError, "Rate limited (429). Retry after #{retry_after}s"
+        end
+
+        if response.code.to_i >= 500
           raise APIError, "HTTP #{response.code}: #{response.body}"
         end
 
@@ -92,13 +118,27 @@ module RerankerRuby
         end
 
         JSON.parse(response.body)
-      rescue APIError => e
+      rescue RateLimitError, APIError => e
         retries += 1
-        if retries <= max_retries && (e.message.include?("429") || e.message.include?("50"))
-          sleep(2 ** (retries - 1))
+        if retries <= @max_retries
+          wait = if e.is_a?(RateLimitError) && e.message =~ /after ([\d.]+)s/
+                   $1.to_f
+                 else
+                   2 ** (retries - 1)
+                 end
+          sleep(wait)
           retry
         end
         raise
+      rescue Net::OpenTimeout, Net::ReadTimeout, Errno::ECONNRESET => e
+        # Reset pooled connection on network errors
+        @conn_mutex.synchronize { @connections.delete("#{uri.host}:#{uri.port}") }
+        retries += 1
+        if retries <= @max_retries
+          sleep(2 ** (retries - 1))
+          retry
+        end
+        raise TimeoutError, "Request failed after #{@max_retries} retries: #{e.message}"
       rescue JSON::ParserError => e
         raise APIError, "Invalid JSON response: #{e.message}"
       end

data/lib/reranker_ruby/cache/memory.rb

@@ -3,6 +3,8 @@
 module RerankerRuby
   module Cache
     class Memory
+      attr_reader :ttl
+
       def initialize(ttl: 3600)
         @ttl = ttl
         @store = {}
@@ -29,6 +31,18 @@ module RerankerRuby
         end
       end
 
+      def invalidate(key)
+        @mutex.synchronize do
+          @store.delete(key)
+        end
+      end
+
+      def invalidate_matching(pattern)
+        @mutex.synchronize do
+          @store.delete_if { |k, _| k.match?(pattern) }
+        end
+      end
+
       def clear
         @mutex.synchronize do
           @store.clear
@@ -40,6 +54,13 @@ module RerankerRuby
           @store.size
         end
       end
+
+      def prune_expired
+        @mutex.synchronize do
+          now = Time.now.to_f
+          @store.delete_if { |_, entry| now - entry[:time] > @ttl }
+        end
+      end
     end
   end
 end

data/lib/reranker_ruby/configuration.rb

@@ -2,14 +2,16 @@
 
 module RerankerRuby
   class Configuration
-    attr_accessor :default_provider, :cohere_api_key, :jina_api_key,
+    attr_accessor :default_provider, :cohere_api_key, :jina_api_key, :voyage_api_key,
                   :default_model, :default_top_k, :cache_store, :cache_ttl,
-                  :logger, :onnx_model, :onnx_model_path, :onnx_cache_dir
+                  :logger, :onnx_model, :onnx_model_path, :onnx_cache_dir,
+                  :timeout, :max_retries
 
     def initialize
       @default_provider = :cohere
       @cohere_api_key = nil
       @jina_api_key = nil
+      @voyage_api_key = nil
       @default_model = nil
       @default_top_k = 10
       @cache_store = nil
@@ -18,6 +20,8 @@ module RerankerRuby
       @onnx_model = nil
       @onnx_model_path = nil
       @onnx_cache_dir = nil
+      @timeout = 30
+      @max_retries = 3
     end
 
     # Build a reranker instance from configuration
@@ -35,6 +39,11 @@ module RerankerRuby
         opts = { api_key: jina_api_key, cache: cache }
         opts[:model] = default_model if default_model
         Jina.new(**opts)
+      when :voyage
+        raise Error, "voyage_api_key is required" unless voyage_api_key
+        opts = { api_key: voyage_api_key, cache: cache, timeout: timeout, max_retries: max_retries }
+        opts[:model] = default_model if default_model
+        Voyage.new(**opts)
       when :onnx
         opts = { cache: cache }
         opts[:model] = onnx_model if onnx_model
@@ -66,10 +75,7 @@ module RerankerRuby
 
   class << self
     def configuration
-      @config_mutex ||= Mutex.new
-      @config_mutex.synchronize do
-        @configuration ||= Configuration.new
-      end
+      @configuration ||= Configuration.new
     end
 
     def configure
@@ -77,19 +83,13 @@ module RerankerRuby
     end
 
     def reset_configuration!
-      @config_mutex ||= Mutex.new
-      @config_mutex.synchronize do
-        @configuration = Configuration.new
-        @reranker = nil
-      end
+      @configuration = Configuration.new
+      @reranker = nil
     end
 
     # Global reranker instance built from configuration
     def reranker
-      @config_mutex ||= Mutex.new
-      @config_mutex.synchronize do
-        @reranker ||= configuration.build_reranker
-      end
+      @reranker ||= configuration.build_reranker
     end
 
     # Convenience method for quick reranking

data/lib/reranker_ruby/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module RerankerRuby
-  VERSION = "0.1.1"
+  VERSION = "0.2.0"
 end

data/lib/reranker_ruby/voyage.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module RerankerRuby
+  class Voyage < Base
+    API_URL = "https://api.voyageai.com/v1/rerank"
+    DEFAULT_MODEL = "rerank-2"
+
+    def initialize(api_key:, model: DEFAULT_MODEL, **options)
+      super(**options)
+      @api_key = api_key
+      @model = model
+    end
+
+    def rerank(query, documents, top_k: 10, model: nil)
+      validate_inputs!(query, documents, top_k)
+      instrument(query: query, document_count: documents.length, top_k: top_k) do
+        with_cache(query, documents, top_k: top_k) do
+          texts = extract_texts(documents)
+
+          response = post(API_URL, {
+            model: model || @model,
+            query: query,
+            documents: texts,
+            top_k: top_k
+          }, headers: {
+            "Authorization" => "Bearer #{@api_key}"
+          })
+
+          response["data"].map do |r|
+            idx = r["index"]
+            Result.new(
+              text: texts[idx],
+              score: r["relevance_score"],
+              index: idx,
+              metadata: extract_metadata(documents[idx])
+            )
+          end.sort
+        end
+      end
+    end
+  end
+end

data/lib/reranker_ruby.rb

@@ -5,6 +5,7 @@ require_relative "reranker_ruby/result"
 require_relative "reranker_ruby/base"
 require_relative "reranker_ruby/cohere"
 require_relative "reranker_ruby/jina"
+require_relative "reranker_ruby/voyage"
 require_relative "reranker_ruby/rrf"
 require_relative "reranker_ruby/model_downloader"
 require_relative "reranker_ruby/score_normalizer"

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: reranker-ruby
 version: !ruby/object:Gem::Version
-  version: 0.1.1
+  version: 0.2.0
 platform: ruby
 authors:
 - Johannes Dwi Cahyo
@@ -94,6 +94,7 @@ files:
 - lib/reranker_ruby/rrf.rb
 - lib/reranker_ruby/score_normalizer.rb
 - lib/reranker_ruby/version.rb
+- lib/reranker_ruby/voyage.rb
 homepage: https://github.com/johannesdwicahyo/reranker-ruby
 licenses:
 - MIT