Class: RagRuby::Embedders::OpenAI

Inherits:
Base
  • Object
show all
Defined in:
lib/rag_ruby/embedders/openai.rb

Constant Summary collapse

ENDPOINT =
"https://api.openai.com/v1/embeddings"
DIMENSIONS =
{
  "text-embedding-3-small" => 1536,
  "text-embedding-3-large" => 3072,
  "text-embedding-ada-002" => 1536
}.freeze

Instance Method Summary collapse

Constructor Details

#initialize(model: "text-embedding-3-small", api_key: nil) ⇒ OpenAI

Returns a new instance of OpenAI.

Raises:

  • (ArgumentError)


18
19
20
21
22
# File 'lib/rag_ruby/embedders/openai.rb', line 18

def initialize(model: "text-embedding-3-small", api_key: nil)
  @model = model
  @api_key = api_key || ENV["OPENAI_API_KEY"]
  raise ArgumentError, "OpenAI API key is required (set OPENAI_API_KEY or pass api_key:)" unless @api_key
end

Instance Method Details

#dimension ⇒ Object



36
37
38
# File 'lib/rag_ruby/embedders/openai.rb', line 36

def dimension
  DIMENSIONS.fetch(@model) { 1536 }
end

#embed(text) ⇒ Object



24
25
26
27
# File 'lib/rag_ruby/embedders/openai.rb', line 24

def embed(text)
  response = request([text])
  response.dig("data", 0, "embedding")
end

#embed_batch(texts) ⇒ Object



29
30
31
32
33
34
# File 'lib/rag_ruby/embedders/openai.rb', line 29

def embed_batch(texts)
  response = request(texts)
  response["data"]
    .sort_by { |d| d["index"] }
    .map { |d| d["embedding"] }
end