From 7736f7b4506392930bfc2e5685f19ae8b4b50628 Mon Sep 17 00:00:00 2001 From: Kevin Sylvestre Date: Mon, 21 Oct 2024 15:40:31 -0700 Subject: [PATCH] Support for generating single / multiple embedding --- Gemfile.lock | 2 +- README.md | 23 ++++++++++++++++---- lib/voyageai/client.rb | 8 ++++++- lib/voyageai/embed.rb | 13 ++++++++---- lib/voyageai/embedding.rb | 37 --------------------------------- lib/voyageai/version.rb | 2 +- spec/factories/embed.rb | 2 +- spec/factories/embedding.rb | 10 --------- spec/voyageai/embed_spec.rb | 30 ++++++++++++++++++++++++-- spec/voyageai/embedding_spec.rb | 28 ------------------------- 10 files changed, 66 insertions(+), 89 deletions(-) delete mode 100644 lib/voyageai/embedding.rb delete mode 100644 spec/factories/embedding.rb delete mode 100644 spec/voyageai/embedding_spec.rb diff --git a/Gemfile.lock b/Gemfile.lock index e080957..7612227 100644 --- a/Gemfile.lock +++ b/Gemfile.lock @@ -1,7 +1,7 @@ PATH remote: . specs: - voyageai (1.0.1) + voyageai (1.1.0) http zeitwerk diff --git a/README.md b/README.md index 76121ae..5509492 100644 --- a/README.md +++ b/README.md @@ -10,6 +10,23 @@ gem install voyageai ## Usage +### Generating Single Embedding + +```ruby +require 'voyageai' + +input = 'A quick brown fox jumps over the lazy dog.' + +voyageai = VoyageAI::Client.new(api_key: 'pa-...') # or configure ENV['VOYAGEAI_API_KEY'] + +embed = voyageai.emed(input) +embed.model # "..." +embed.usage # "#" +embed.embedding # [0.0, ...] +``` + +### Generating Multiple Embeddings + ```ruby require 'voyageai' @@ -22,10 +39,8 @@ input = [ voyageai = VoyageAI::Client.new(api_key: 'pa-...') # or configure ENV['VOYAGEAI_API_KEY'] -result = voyageai.embed(input) +embed = voyageai.embed(input) embed.model # "..." embed.usage # "#" -embed.embeddings.each do |embedding| - embedding.index # "# -end +embed.embeddings # [[0.0, ...], ...] ``` diff --git a/lib/voyageai/client.rb b/lib/voyageai/client.rb index 4fa791d..8174aaa 100644 --- a/lib/voyageai/client.rb +++ b/lib/voyageai/client.rb @@ -33,7 +33,7 @@ def inspect # # @return [Embedding] def embed(input, model: Model::VOYAGE) - payload = { input: input, model: model } + payload = { input: arrayify(input), model: model } response = HTTP .accept(:json) .auth("Bearer #{@api_key}") @@ -43,5 +43,11 @@ def embed(input, model: Model::VOYAGE) Embed.parse(data: response.parse) end + + private + + def arrayify(input) + input.is_a?(Array) ? input : [input] + end end end diff --git a/lib/voyageai/embed.rb b/lib/voyageai/embed.rb index 0fe738f..ba5ffca 100644 --- a/lib/voyageai/embed.rb +++ b/lib/voyageai/embed.rb @@ -15,7 +15,7 @@ class Embed attr_accessor :usage # @!attribute [rw] embeddings - # @return [Array] + # @return [Array>] attr_accessor :embeddings # @param data [Hash] @@ -23,9 +23,7 @@ class Embed def self.parse(data:) model = data["model"] usage = Usage.parse(data: data["usage"]) - embeddings = data["data"].map do |embedding_data| - Embedding.parse(data: embedding_data) - end + embeddings = data["data"].map { |embedding_data| embedding_data["embedding"] } Embed.new(model: model, usage: usage, embeddings: embeddings) end @@ -43,5 +41,12 @@ def initialize(model:, usage:, embeddings:) def inspect "#<#{self.class.name} model=#{@model.inspect} embeddings=#{@embeddings.inspect} usage=#{@usage.inspect}>" end + + # @param index [Integer] optional + # + # @return [Array] + def embedding(index: 0) + @embeddings[index] + end end end diff --git a/lib/voyageai/embedding.rb b/lib/voyageai/embedding.rb deleted file mode 100644 index ec6e057..0000000 --- a/lib/voyageai/embedding.rb +++ /dev/null @@ -1,37 +0,0 @@ -# frozen_string_literal: true - -module VoyageAI - # An embedding returned by the VoyageAI API. - # - # @example - # VoyageAI::Embedding.new(index: 0, embedding: [0.0, 1.0, 2.0, 3.0]) - class Embedding - # @!attribute [rw] model - # @return [Integer] - attr_accessor :index - - # @!attribute [rw] embedding - # @return [Array] - attr_accessor :embedding - - # @param data [Hash] - # @return [Embedding] - def self.parse(data:) - index = data["index"] - embedding = data["embedding"] - new(index:, embedding:) - end - - # @param index [Integer] - # @param embedding [Array] - def initialize(index:, embedding:) - @index = index - @embedding = embedding - end - - # @return [String] - def inspect - "#<#{self.class.name} index=#{index} embedding=#{embedding.inspect}>" - end - end -end diff --git a/lib/voyageai/version.rb b/lib/voyageai/version.rb index 6418524..50a8130 100644 --- a/lib/voyageai/version.rb +++ b/lib/voyageai/version.rb @@ -1,5 +1,5 @@ # frozen_string_literal: true module VoyageAI - VERSION = "1.0.1" + VERSION = "1.1.0" end diff --git a/spec/factories/embed.rb b/spec/factories/embed.rb index 3f0571b..1cf0072 100644 --- a/spec/factories/embed.rb +++ b/spec/factories/embed.rb @@ -6,6 +6,6 @@ model { VoyageAI::Model::VOYAGE } usage - embeddings { [] } + embeddings { [[0.0]] } end end diff --git a/spec/factories/embedding.rb b/spec/factories/embedding.rb deleted file mode 100644 index 13549d0..0000000 --- a/spec/factories/embedding.rb +++ /dev/null @@ -1,10 +0,0 @@ -# frozen_string_literal: true - -FactoryBot.define do - factory :embedding, class: "VoyageAI::Embedding" do - initialize_with { new(index:, embedding:) } - - index { 0 } - embedding { [] } - end -end diff --git a/spec/voyageai/embed_spec.rb b/spec/voyageai/embed_spec.rb index 412364e..2451e85 100644 --- a/spec/voyageai/embed_spec.rb +++ b/spec/voyageai/embed_spec.rb @@ -10,7 +10,7 @@ "data" => [ { "object" => "embedding", - "embedding" => [0.0, 1.0, 2.0, 3.0], + "embedding" => [0.0], "index" => 0, }, ], @@ -22,15 +22,41 @@ it "parses a hash and returns an instance" do expect(parse).to be_a(described_class) end + + it "sets embeddings to an array of arrays" do + expect(parse.embeddings).to eql [[0.0]] + end end describe "#inspect" do subject(:inspect) { embed.inspect } - let(:embed) { build(:embed) } + let(:embed) { build(:embed, embeddings: []) } it "returns a string" do expect(inspect).to eql '#>' end end + + describe "#embedding" do + subject(:embedding) { embed.embedding(index: index) } + + let(:embed) { build(:embed) } + + context "with an index that exists" do + let(:index) { 0 } + + it "returns an array" do + expect(embedding).to eql [0.0] + end + end + + context "when the index that does not exist" do + let(:index) { 2 } + + it "returns nil" do + expect(embedding).to be_nil + end + end + end end diff --git a/spec/voyageai/embedding_spec.rb b/spec/voyageai/embedding_spec.rb deleted file mode 100644 index 48980d4..0000000 --- a/spec/voyageai/embedding_spec.rb +++ /dev/null @@ -1,28 +0,0 @@ -# frozen_string_literal: true - -RSpec.describe VoyageAI::Embedding do - describe ".parse" do - subject(:parse) { described_class.parse(data: data) } - - let(:data) do - { - "embedding" => [], - "index" => 0, - } - end - - it "parses a hash and returns an instance" do - expect(parse).to be_a(described_class) - end - end - - describe "#inspect" do - subject(:inspect) { embedding.inspect } - - let(:embedding) { build(:embedding) } - - it "returns a string" do - expect(inspect).to eql "#" - end - end -end