voyageai 1.2.0 → 1.5.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 8d127480aae597e6c15b398061d16beded7be8961fd433372b2550ea1b643d37
-  data.tar.gz: 0a958352786d2994de9bdfbbb198850e15d245a9aed6971efae0f30518494878
+  metadata.gz: 651f676fcb2fa068cf44f37fda02d1185b776ca83fc0f6099362d5e4deb7a250
+  data.tar.gz: f3a22402be82fd02ec94628aeab067aec45cfcc6c203f7d75af13a00f95b580b
 SHA512:
-  metadata.gz: 346dc805e9a46813198189b83805a6d2718cb1a4c244af08f8a1fb5cc4cc251f0fd7ff1b06104669b588da56b5d8fcda3f06795a64fac5a2ee7dc6e8e86faa1f
-  data.tar.gz: 9cf273cdbf521f0b13fc9b30a58517ae7ebdfce1cb23801e64630bcf4887f17d4ed61058d0eb72533fe801e221a7a1fe91df84b802cb5266d7a405af7b68981a
+  metadata.gz: 0b70798b6ee15febae3f263def8f118b4177f73e553bf3463e77f2638259b62d95025dc6336049eedec71fdccdac80512e54c55bd65b70333fc3f29518d97367
+  data.tar.gz: 2aa09dc2ea94771f309af4e40aa4340842b1181e1e82ddc270b75bddd5120044d0dd195bf687e302cf270da12c312261dc6a523e5b69a0184bacf0aafbb10da1

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
@@ -58,3 +82,81 @@ VoyageAI.configure do |config|
   config.logger = Logger.new(STDOUT)
 end
 ```
+
+## Examples
+
+```ruby
+require "voyageai"
+
+Entry = Data.define(:document, :embedding)
+
+VOYAGEAI = VoyageAI::Client.new
+
+DOCUMENTS = [
+  "John is a musician.",
+  "Paul is a plumber.",
+  "George is a teacher.",
+  "Ringo is a doctor.",
+  "Lisa is a lawyer.",
+  "Stuart is a painter.",
+  "Brian is a writer.",
+  "Jane is a chef.",
+  "Bill is a nurse.",
+  "Susan is a carpenter.",
+].freeze
+
+embeddings = VOYAGEAI.embed(DOCUMENTS, input_type: "document").embeddings
+
+ENTRIES = DOCUMENTS.zip(embeddings).map do |document, embedding|
+  Entry.new(document:, embedding:)
+end
+
+# @param src [Array<Float>]
+# @param dst [Array<Float>]
+#
+# @return [Float]
+def euclidean_distance(src, dst)
+  Math.sqrt(src.zip(dst).map { |a, b| (a - b)**2 }.reduce(:+))
+end
+
+# @param query [String]
+# @param limit [Integer]
+#
+# @return [Array<String>]
+def nearest_documents(query:, limit: 4)
+  embedding = VOYAGEAI.embed(query, input_type: "query").embedding
+
+  ENTRIES
+    .sort_by { |entry| euclidean_distance(entry.embedding, embedding) }
+    .first(limit)
+    .map(&:document)
+end
+
+# @param query [String]
+def search(query:)
+  documents = nearest_documents(query:)
+
+  results = VOYAGEAI.rerank(query:, documents:, top_k: 2).results
+
+  puts "query=#{query.inspect}"
+  results.each do |reranking|
+    document = documents[reranking.index]
+    puts("document=#{document.inspect} relevance_score=#{reranking.relevance_score}")
+  end
+end
+
+search(query: "What do George and Ringo do?")
+search(query: "Who works in the medical field?")
+```
+
+```
+query="What do George and Ringo do?"
+document="Ringo is a doctor." relevance_score=0.67968755
+document="George is a teacher." relevance_score=0.58593755
+```
+
+```
+query="Who works in the medical field?"
+document="Bill is a nurse." relevance_score=0.55078125
+document="Ringo is a doctor." relevance_score=0.50390625
+```

data/lib/voyageai/client.rb

@@ -37,12 +37,14 @@ module VoyageAI
       "#<#{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")
     # @param input [String, Array<String>] required (e.g. "Sample" or ["Sample 1", "Sample 2", ...])
+    # @param model [String] optional (e.g. VoyageAI::Model::VOYAGE or "voyage-3")
+    # @param input_type [Symbol] optional (e.g. :query or :document)
+    # @param truncation [Boolean] optional
     #
     # @return [Embedding]
-    def embed(input, model: Model::VOYAGE)
-      payload = { input: arrayify(input), model: model }
+    def embed(input, model: Model::VOYAGE, input_type: nil, truncation: nil)
+      payload = { input: arrayify(input), model: model, truncation:, input_type: }.compact
       response = http.accept(:json).post("/#{@version}/embeddings", json: payload)
 
       raise RequestError.new(response:) unless response.status.ok?
@@ -50,6 +52,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

@@ -3,17 +3,24 @@
 module VoyageAI
   module Model
     VOYAGE_3 = "voyage-3"
+    VOYAGE_3_LARGE = "voyage-3-large"
     VOYAGE_3_LITE = "voyage-3-lite"
     VOYAGE_FINANCE_2 = "voyage-finance-2"
     VOYAGE_MULTILINGUAL_2 = "age-multilingual-2"
     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.5.0"
 end

metadata

@@ -1,14 +1,13 @@
 --- !ruby/object:Gem::Specification
 name: voyageai
 version: !ruby/object:Gem::Version
-  version: 1.2.0
+  version: 1.5.0
 platform: ruby
 authors:
 - Kevin Sylvestre
-autorequire: 
 bindir: bin
 cert_chain: []
-date: 2024-10-22 00:00:00.000000000 Z
+date: 2025-01-08 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: http
@@ -55,6 +54,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
@@ -64,7 +65,6 @@ metadata:
   homepage_uri: https://github.com/ksylvest/voyageai
   changelog_uri: https://github.com/ksylvest/voyageai/releases
   rubygems_mfa_required: 'true'
-post_install_message: 
 rdoc_options: []
 require_paths:
 - lib
@@ -79,8 +79,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.5.18
-signing_key: 
+rubygems_version: 3.6.2
 specification_version: 4
 summary: A client for voyageai.com.
 test_files: []