voyageai 1.1.1 → 1.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: c834adfa635ba86d940fab56c4ac8bb46086fc6a025319d392c818533389e7dd
-  data.tar.gz: a610e60874f1440eaf05ae2f545350b79fd432b786ad60392251032e5b1cc7e1
+  metadata.gz: ac2d6a908abb5d1383cdf74f736b5aa7eae9fc9a66634ec1fd40ce8710571763
+  data.tar.gz: 7df317f435acc60bde69219f7dc50e19952433dce8d4dee9ebda10a8a6b0ba46
 SHA512:
-  metadata.gz: f9560b027aa73f921175a2b06ff6f551ed582f945b8838dd8fdd632ee2d321be89312fad8d3d0179490dbc06e1342a1b6f6973374953c797ce00b2a43f4a334b
-  data.tar.gz: efb6b2f84bfe349d2f909895d14c9291ace8c7af16179d4e1b83fc911b2216f69aa71826a7996fcdd97fa1b347f113eb976a8ca941ab2643726a6e20fa1b1aa6
+  metadata.gz: af23bbd90f54d69d912223071a4a9ca3ba9d4128c10e7d5e60ce57632c54a244f4dc656ae6dd3a923d8b8273e1836ae670401a8de7de155350d1a1a5e2e24a1b
+  data.tar.gz: a778faecde1b31e4c7a9c8955019af8ec618076dff925b3d26dce0bceeaae054174f24cfeb120d4c076750fcc730679338014c50c5ffdfe47ee2bad9a10700a4

data/README.md

@@ -10,7 +10,9 @@ gem install voyageai
 
 ## Usage
 
-### Generating Single Embedding
+### Embedding
+
+#### Generating Single Embedding
 
 ```ruby
 require 'voyageai'
@@ -25,7 +27,7 @@ embed.usage # "#<VoyageAI::Usage total_tokens=...>"
 embed.embedding # [0.0, ...]
 ```
 
-### Generating Multiple Embeddings
+#### Generating Multiple Embeddings
 
 ```ruby
 require 'voyageai'
@@ -44,3 +46,39 @@ embed.model # "..."
 embed.usage # "#<VoyageAI::Usage total_tokens=...>"
 embed.embeddings # [[0.0, ...], ...]
 ```
+
+### Reranking
+
+```ruby
+require 'voyageai'
+
+query = 'Who is the best person to call for a toilet?'
+
+documents = [
+  'John is a musician.',
+  'Paul is a plumber.',
+  'George is a teacher.',
+  'Ringo is a doctor.',
+]
+
+voyageai = VoyageAI::Client.new(api_key: 'pa-...') # or configure ENV['VOYAGEAI_API_KEY']
+
+rerank = voyageai.rerank(query:, documents:, top_k: 3)
+rerank.model # "..."
+rerank.usage # "#<VoyageAI::Usage total_tokens=...>"
+rerank.results # [#<VoyageAI::Reranking index=0 relevance_score=0.5>]
+```
+
+## Configuration
+
+```ruby
+require 'voyageai'
+
+VoyageAI.configure do |config|
+  config.api_key = 'pa-...' # defaults to ENV['VOYAGEAI_API_KEY']
+  config.host = 'https://api.voyageai.com'
+  config.version = 'v1'
+  config.timeout = 15 # seconds
+  config.logger = Logger.new(STDOUT)
+end
+```

data/lib/voyageai/client.rb

@@ -3,8 +3,6 @@
 module VoyageAI
   # A client for interacting with the VoyageAI API.
   class Client
-    BASE_URL = "https://api.voyageai.com/v1"
-
     # An error raised for any HTTP issues.
     class RequestError < StandardError
       attr_accessor :response
@@ -17,15 +15,26 @@ module VoyageAI
     end
 
     # @param api_key [String] optional if ENV['VOYAGEAI_API_KEY'] otherwise required
-    def initialize(api_key: ENV.fetch("VOYAGEAI_API_KEY", nil))
+    # @param host [String] optional
+    # @param version [String] optional
+    # @param logger [Logger] optional
+    def initialize(
+      api_key: VoyageAI.config.api_key,
+      host: VoyageAI.config.host,
+      version: VoyageAI.config.version,
+      logger: VoyageAI.config.logger
+    )
       @api_key = api_key || raise(ArgumentError, "api_key is required or ENV['VOYAGEAI_API_KEY'] must be present")
+      @host = host
+      @version = version
+      @logger = logger
     end
 
     # @return [String]
     def inspect
       masked_api_key = "#{@api_key[..4]}***"
 
-      "#<#{self.class.name} api_key=#{masked_api_key.inspect}>"
+      "#<#{self.class.name} api_key=#{masked_api_key.inspect} host=#{@host.inspect} version=#{@version.inspect}>"
     end
 
     # @param model [String] optional (e.g. VoyageAI::Model::VOYAGE or "voyage-3")
@@ -34,18 +43,39 @@ module VoyageAI
     # @return [Embedding]
     def embed(input, model: Model::VOYAGE)
       payload = { input: arrayify(input), model: model }
-      response = HTTP
-        .accept(:json)
-        .auth("Bearer #{@api_key}")
-        .post("#{BASE_URL}/embeddings", json: payload)
+      response = http.accept(:json).post("/#{@version}/embeddings", json: payload)
 
       raise RequestError.new(response:) unless response.status.ok?
 
       Embed.parse(data: response.parse)
     end
 
+    # @param query [String] required
+    # @param documents [Array<String>] required
+    # @param model [String] optional (e.g. VoyageAI::Model::RERANK or "rerank-2")
+    # @param top_k [Integer] optional
+    # @param truncation [Boolean] optional
+    def rerank(query:, documents:, model: Model::RERANK, top_k: nil, truncation: nil)
+      payload = { query:, documents:, model:, top_k:, truncation: }.compact
+      response = http.accept(:json).post("/#{@version}/rerank", json: payload)
+
+      raise RequestError.new(response:) unless response.status.ok?
+
+      Rerank.parse(data: response.parse)
+    end
+
     private
 
+    # @return [HTTP::Client]
+    def http
+      @http ||= begin
+        http = HTTP.auth("Bearer #{@api_key}").persistent(@host)
+        http = http.use(instrumentation: { instrumenter: Instrumentation.new(logger: @logger) }) if @logger
+        http = http.timeout(@timeout) if @timeout
+        http
+      end
+    end
+
     def arrayify(input)
       input.is_a?(Array) ? input : [input]
     end

data/lib/voyageai/config.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+module VoyageAI
+  # A configuration for VoyageAI.
+  #
+  # @example:
+  #   VoyageAI.configure do |config|
+  #     config.api_key = '...'
+  #     config.host = 'https://api.voyageai.com'
+  #     config.version = 'v1'
+  #     config.logger = Logger.new(STDOUT)
+  #     config.timeout = 15
+  #   end
+  class Config
+    DEFAULT_HOST = "https://api.voyageai.com"
+    DEFAULT_VERSION = "v1"
+
+    # @!attribute [rw] api_key
+    #   @return [String, nil]
+    attr_accessor :api_key
+
+    # @!attribute [rw] host
+    #   @return [String, nil]
+    attr_accessor :host
+
+    # @!attribute [rw] version
+    #   @return [String, nil]
+    attr_accessor :version
+
+    # @!attribute [rw] logger
+    #   @return [Logger, nil]
+    attr_accessor :logger
+
+    # @!attribute [rw] timeout
+    #   @return [Integer, nil]
+    attr_accessor :timeout
+
+    # @param api_key [String] optional - defaults to `ENV['VOYAGE_API_KEY']`
+    # @param host [String] optional - defaults to `ENV['VOYAGEAI_HOST']`` w/ fallback to `https://api.voyageai.com`
+    # @param version [String] optional - defaults to `ENV['DEFAULT_VERSION']` w/ fallback to `v1`
+    # @param logger [Logger] optional
+    # @param timeout [Integer] optional
+    def initialize(
+      api_key: ENV.fetch("VOYAGEAI_API_KEY", nil),
+      host: ENV.fetch("VOYAGEAI_HOST", DEFAULT_HOST),
+      version: ENV.fetch("VOYAGEAI_VERSION", DEFAULT_VERSION),
+      logger: nil,
+      timeout: nil
+    )
+      @api_key = api_key
+      @host = host
+      @version = version
+      @logger = logger
+      @timeout = timeout
+    end
+  end
+end

data/lib/voyageai/embed.rb

@@ -25,7 +25,7 @@ module VoyageAI
       usage = Usage.parse(data: data["usage"])
       embeddings = data["data"].map { |embedding_data| embedding_data["embedding"] }
 
-      Embed.new(model: model, usage: usage, embeddings: embeddings)
+      new(model: model, usage: usage, embeddings: embeddings)
     end
 
     # @param model [String]

data/lib/voyageai/instrumentation.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module VoyageAI
+  # Used for logging via HTTP.
+  class Instrumentation
+    # @param logger [Logger]
+    def initialize(logger:)
+      @logger = logger
+    end
+
+    # @param name [String]
+    # @param payload [Hash]
+    # @option payload [Exception] :error
+    def instrument(name, payload = {})
+      error = payload[:error]
+      return unless error
+
+      @logger.error("#{name}: #{error.message}")
+    end
+
+    # @param payload [Hash]
+    # @option payload [HTTP::Request] :request
+    def start(_, payload)
+      request = payload[:request]
+      @logger.info("#{request.verb.upcase} #{request.uri}")
+    end
+
+    # @param payload [Hash]
+    # @option payload [HTTP::Response] :response
+    def finish(_, payload)
+      response = payload[:response]
+      @logger.info("#{response.status.code} #{response.status.reason}")
+    end
+  end
+end

data/lib/voyageai/model.rb

@@ -9,11 +9,17 @@ module VoyageAI
     VOYAGE_LAW_2 = "voyage-law-2"
     VOYAGE_CODE_2 = "voyage-code-2"
 
+    RERANK_2 = "rerank-2"
+    RERANK_2_LITE = "rerank-2-lite"
+
     VOYAGE = VOYAGE_3
     VOYAGE_LITE = VOYAGE_3_LITE
     VOYAGE_FINANCE = VOYAGE_FINANCE_2
     VOYAGE_MULTILINGUAL = VOYAGE_MULTILINGUAL_2
     VOYAGE_LAW = VOYAGE_LAW_2
     VOYAGE_CODE = VOYAGE_CODE_2
+
+    RERANK = RERANK_2
+    RERANK_LITE = RERANK_2_LITE
   end
 end

data/lib/voyageai/rerank.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+module VoyageAI
+  # The response for an rerank request that wraps the model / usage / results.
+  #
+  # @example
+  #   VoyageAI::Rerank.new(model: "rerank-2", usage: VoyageAI::Usage.new(total_tokens: 0), results: [])
+  class Rerank
+    # @!attribute [rw] model
+    #   @return [String]
+    attr_accessor :model
+
+    # @!attribute [rw] usage
+    #   @return [Usage]
+    attr_accessor :usage
+
+    # @!attribute [rw] reranking
+    #   @return [Array<Array<Float>>]
+    attr_accessor :results
+
+    # @param data [Hash]
+    # @return [Embed]
+    def self.parse(data:)
+      model = data["model"]
+      usage = Usage.parse(data: data["usage"])
+
+      results = data["data"].map { |entry| Reranking.parse(data: entry) }
+
+      new(model:, usage:, results:)
+    end
+
+    # @param model [String]
+    # @param usage [Usage]
+    # @param results [Array<Reranking>]
+    def initialize(model:, usage:, results:)
+      @model = model
+      @usage = usage
+      @results = results
+    end
+
+    # @return [String]
+    def inspect
+      "#<#{self.class.name} model=#{@model.inspect} usage=#{@usage.inspect}>"
+    end
+  end
+end

data/lib/voyageai/reranking.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+module VoyageAI
+  # The response of an individual re-ranking.
+  #
+  # @example
+  #   VoyageAI::Reranking.new(index: 0, document: "Sample", relevance_score: 0.0)
+  class Reranking
+    # @!attribute [rw] model
+    #   @return [Integer]
+    attr_accessor :index
+
+    # @!attribute [rw] document
+    #   @return [String]
+    attr_accessor :document
+
+    # @!attribute [rw] relevance_score
+    #   @return [Float]
+    attr_accessor :relevance_score
+
+    # @param index [Integer]
+    # @param document [String]
+    # @param relevance_score [Float]
+    def initialize(index:, document:, relevance_score:)
+      @index = index
+      @document = document
+      @relevance_score = relevance_score
+    end
+
+    # @return [String]
+    def inspect
+      "#<#{self.class.name} index=#{@index} relevance_score=#{@relevance_score}>"
+    end
+
+    # @param data [Hash]
+    #
+    # @return [Reranking]
+    def self.parse(data:)
+      index = data["index"]
+      document = data["document"]
+      relevance_score = data["relevance_score"]
+
+      new(index: index, document: document, relevance_score: relevance_score)
+    end
+  end
+end

data/lib/voyageai/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module VoyageAI
-  VERSION = "1.1.1"
+  VERSION = "1.3.0"
 end

data/lib/voyageai.rb

@@ -7,6 +7,15 @@ loader = Zeitwerk::Loader.for_gem
 loader.inflector.inflect "voyageai" => "VoyageAI"
 loader.setup
 
+# The main entrypoint for VoyageAI.
 module VoyageAI
-  class Error < StandardError; end
+  # @return [VoyageAI::Config]
+  def self.config
+    @config ||= Config.new
+  end
+
+  # @yield [VoyageAI::Config]
+  def self.configure
+    yield config
+  end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: voyageai
 version: !ruby/object:Gem::Version
-  version: 1.1.1
+  version: 1.3.0
 platform: ruby
 authors:
 - Kevin Sylvestre
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2024-10-21 00:00:00.000000000 Z
+date: 2024-10-22 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: http
@@ -51,8 +51,12 @@ files:
 - bin/setup
 - lib/voyageai.rb
 - lib/voyageai/client.rb
+- lib/voyageai/config.rb
 - lib/voyageai/embed.rb
+- lib/voyageai/instrumentation.rb
 - lib/voyageai/model.rb
+- lib/voyageai/rerank.rb
+- lib/voyageai/reranking.rb
 - lib/voyageai/usage.rb
 - lib/voyageai/version.rb
 homepage: https://github.com/ksylvest/voyageai