Class: Gemini::Embeddings

Inherits:
Object
  • Object
show all
Defined in:
lib/gemini/embeddings.rb

Constant Summary collapse

DEFAULT_MODEL = 
"gemini-embedding-001".freeze
VALID_TASK_TYPES = 
%w[
  RETRIEVAL_QUERY
  RETRIEVAL_DOCUMENT
  SEMANTIC_SIMILARITY
  CLASSIFICATION
  CLUSTERING
  QUESTION_ANSWERING
  FACT_VERIFICATION
  CODE_RETRIEVAL_QUERY
].freeze

Instance Method Summary collapse

Constructor Details

#initialize(client:) ⇒ Embeddings

Returns a new instance of Embeddings.



16
17
18
 
# File 'lib/gemini/embeddings.rb', line 16

def initialize(client:)
  @client = client
end

Instance Method Details

#batch_create(inputs:, model: DEFAULT_MODEL, task_type: nil, title: nil, output_dimensionality: nil, **parameters) ⇒ Object

Generate embeddings for multiple inputs in a single batch request



49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
 
# File 'lib/gemini/embeddings.rb', line 49

def batch_create(inputs:, model: DEFAULT_MODEL, task_type: nil, title: nil,
                 output_dimensionality: nil, **parameters)
  requests = inputs.map do |input|
    req = build_embed_payload(
      input: input,
      task_type: task_type,
      title: title,
      output_dimensionality: output_dimensionality
    )
    req[:model] = "models/#{normalize_model(model)}"
    req
  end

  payload = { requests: requests }.merge(parameters)

  response = @client.json_post(
    path: "models/#{normalize_model(model)}:batchEmbedContents",
    parameters: payload
  )
  Gemini::Response.new(response)
end

#create(input:, model: DEFAULT_MODEL, task_type: nil, title: nil, output_dimensionality: nil, **parameters) ⇒ Object

Generate an embedding for a single content, or batch when input is an Array



21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
 
# File 'lib/gemini/embeddings.rb', line 21

def create(input:, model: DEFAULT_MODEL, task_type: nil, title: nil,
           output_dimensionality: nil, **parameters)
  if input.is_a?(Array)
    return batch_create(
      inputs: input,
      model: model,
      task_type: task_type,
      title: title,
      output_dimensionality: output_dimensionality,
      **parameters
    )
  end

  payload = build_embed_payload(
    input: input,
    task_type: task_type,
    title: title,
    output_dimensionality: output_dimensionality
  ).merge(parameters)

  response = @client.json_post(
    path: "models/#{normalize_model(model)}:embedContent",
    parameters: payload
  )
  Gemini::Response.new(response)
end