voyageai 1.2.0 → 1.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 8d127480aae597e6c15b398061d16beded7be8961fd433372b2550ea1b643d37
-  data.tar.gz: 0a958352786d2994de9bdfbbb198850e15d245a9aed6971efae0f30518494878
+  metadata.gz: ac2d6a908abb5d1383cdf74f736b5aa7eae9fc9a66634ec1fd40ce8710571763
+  data.tar.gz: 7df317f435acc60bde69219f7dc50e19952433dce8d4dee9ebda10a8a6b0ba46
 SHA512:
-  metadata.gz: 346dc805e9a46813198189b83805a6d2718cb1a4c244af08f8a1fb5cc4cc251f0fd7ff1b06104669b588da56b5d8fcda3f06795a64fac5a2ee7dc6e8e86faa1f
-  data.tar.gz: 9cf273cdbf521f0b13fc9b30a58517ae7ebdfce1cb23801e64630bcf4887f17d4ed61058d0eb72533fe801e221a7a1fe91df84b802cb5266d7a405af7b68981a
+  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'
@@ -45,6 +47,28 @@ 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

data/lib/voyageai/client.rb

@@ -50,6 +50,20 @@ module VoyageAI
       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]

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/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.2.0"
+  VERSION = "1.3.0"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: voyageai
 version: !ruby/object:Gem::Version
-  version: 1.2.0
+  version: 1.3.0
 platform: ruby
 authors:
 - Kevin Sylvestre
@@ -55,6 +55,8 @@ files:
 - 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