Module: Rogalik::Plugins::OpenAIEmbeddings::EmbeddingsMethods

Defined in:
lib/rogalik/plugins/open_ai_embeddings.rb

Instance Attribute Summary collapse

Instance Method Summary collapse

Instance Attribute Details

#api_key ⇒ String (readonly)

Returns OpenAI API key.

Returns:

  • (String) —

    OpenAI API key



11
12
13
 
# File 'lib/rogalik/plugins/open_ai_embeddings.rb', line 11

def api_key
  @api_key
end

#model ⇒ String (readonly)

Returns embedding model name.

Returns:

  • (String) —

    embedding model name



14
15
16
 
# File 'lib/rogalik/plugins/open_ai_embeddings.rb', line 14

def model
  @model
end

Instance Method Details

#client ⇒ OpenAI::Client

Returns a memoized OpenAI client.

Returns:

  • (OpenAI::Client) 


29
30
31
 
# File 'lib/rogalik/plugins/open_ai_embeddings.rb', line 29

def client
  @client ||= OpenAI::Client.new(api_key: api_key)
end

#embed_documents(documents) ⇒ Array<Pipeline::Document>

Generates embeddings for each document and assigns them in place.

Parameters:

  • documents (Array<Pipeline::Document>) —

    documents whose page_content will be embedded

Returns:

  • (Array<Pipeline::Document>) —

    the same documents with embeddings populated



37
38
39
40
41
42
43
44
45
46
 
# File 'lib/rogalik/plugins/open_ai_embeddings.rb', line 37

def embed_documents(documents)
  response = client.embeddings.create(input: documents.map(&:page_content), model: model)
  embeddings = response.data

  documents.each_with_index do |doc, index|
    doc.embeddings = embeddings[index].embedding
  end

  documents
end

#embed_text(text) ⇒ Array<Float>

Generates an embedding vector for a single string.

Parameters:

  • text (String) —

    the text to embed

Returns:

  • (Array<Float>) —

    the embedding vector



52
53
54
55
 
# File 'lib/rogalik/plugins/open_ai_embeddings.rb', line 52

def embed_text(text)
  response = client.embeddings.create(input: text, model: model)
  response.data.first.embedding
end

#initialize(config = Rogalik.configuration.embeddings) ⇒ Object

Parameters:

Raises:

  • (ArgumentError) 


17
18
19
20
21
22
23
24
 
# File 'lib/rogalik/plugins/open_ai_embeddings.rb', line 17

def initialize(config = Rogalik.configuration.embeddings)
  api_key = config.api_key || ENV["OPENAI_API_KEY"]
  raise ArgumentError, "OpenAI API key is required, set through env OPENAI_API_KEY or config" if api_key.nil?
  raise ArgumentError, "OpenAI model name is required" if config.model.nil?

  @api_key = api_key
  @model = config.model
end