Support for generating single / multiple embedding

Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    voyageai (1.0.1)
+    voyageai (1.1.0)
       http
       zeitwerk
 

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 # "#<VoyageAI::Usage total_tokens=...>"
+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 # "#<VoyageAI::Usage total_tokens=...>"
-embed.embeddings.each do |embedding|
-  embedding.index # "#<VoyageAI::Embedding index=... embedding=...>
-end
+embed.embeddings # [[0.0, ...], ...]
 ```

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

lib/voyageai/embed.rb

@@ -15,17 +15,15 @@ class Embed
     attr_accessor :usage
 
     # @!attribute [rw] embeddings
-    #   @return [Array<Embedding>]
+    #   @return [Array<Array<Float>>]
     attr_accessor :embeddings
 
     # @param data [Hash]
     # @return [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<Float>]
+    def embedding(index: 0)
+      @embeddings[index]
+    end
   end
 end

lib/voyageai/version.rb

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

spec/factories/embed.rb

@@ -6,6 +6,6 @@
 
     model { VoyageAI::Model::VOYAGE }
     usage
-    embeddings { [] }
+    embeddings { [[0.0]] }
   end
 end

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 '#<VoyageAI::Embed model="voyage-3" embeddings=[] usage=#<VoyageAI::Usage total_tokens=0>>'
     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