Sha256: d659f34f1f6549eceb019093cf6dddcf3ba8be5cb01824bd9a01caada91f325a

Contents?: true

Size: 1.21 KB

Versions: 1

Compression:

Stored size: 1.21 KB

Contents

# frozen_string_literal: true

require 'openai'

module Boxcars
  module VectorStores
    class EmbedViaOpenAI
      include VectorStore

      attr_accessor :texts, :client, :model

      def initialize(texts:, client:, model: 'text-embedding-ada-002')
        validate_params(texts, client)
        @texts = texts
        @client = client
        @model = model
      end

      def call
        texts.map do |text|
          embedding = embedding_with_retry(model: model, input: strip_new_lines(text))
          {
            embedding: embedding,
            dim: embedding.size
          }
        end
      end

      private

      def validate_params(texts, client)
        raise_error 'texts must be an array of strings' unless texts.is_a?(Array) && texts.all? { |text| text.is_a?(String) }
        raise_error 'openai_connection must be an OpenAI::Client' unless client.is_a?(OpenAI::Client)
      end

      def embedding_with_retry(request)
        response = @client.embeddings(parameters: request)
        response['data'][0]['embedding']
      end

      def strip_new_lines(text)
        text.gsub("\n", ' ')
      end

      def raise_error(message)
        raise ::Boxcars::ValueError, message
      end
    end
  end
end

Version data entries

1 entries across 1 versions & 1 rubygems

Version Path
boxcars-0.2.10 lib/boxcars/boxcar/vector_stores/embed_via_open_ai.rb