Class: ActiveRecordVector::Providers::Cohere

Inherits:
Base
  • Object
show all
Defined in:
lib/active_record_vector/providers/cohere.rb

Overview

Cohere Embedding Provider

Constant Summary collapse

DEFAULT_MODEL = 
"embed-english-v3.0"
API_ENDPOINT = 
"https://api.cohere.com/v1/embed"

Instance Attribute Summary

Attributes inherited from Base

#options

Instance Method Summary collapse

Methods inherited from Base

#embed_batch, #initialize

Constructor Details

This class inherits a constructor from ActiveRecordVector::Providers::Base

Instance Method Details

#embed(text) ⇒ Object

Raises:



14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
 
# File 'lib/active_record_vector/providers/cohere.rb', line 14

def embed(text)
  api_key = options[:api_key] || ENV.fetch("COHERE_API_KEY", nil)
  raise Error, "Cohere API key missing. Set ENV['COHERE_API_KEY'] or pass api_key: '...'" unless api_key

  uri = URI.parse(options[:endpoint] || API_ENDPOINT)
  http = Net::HTTP.new(uri.host, uri.port)
  http.use_ssl = (uri.scheme == "https")

  payload = {
    texts: [text],
    model: options[:model] || DEFAULT_MODEL,
    input_type: options[:input_type] || "search_document"
  }

  request = Net::HTTP::Post.new(uri.request_uri, {
                                  "Content-Type" => "application/json",
                                  "Authorization" => "Bearer #{api_key}"
                                })
  request.body = payload.to_json

  response = http.request(request)
  raise Error, "Cohere API Error (#{response.code}): #{response.body}" unless response.is_a?(Net::HTTPSuccess)

  data = JSON.parse(response.body)
  data.dig("embeddings", 0) || []
end