Class: Rager::Rerank::Providers::Cohere

Inherits:
Abstract
  • Object
show all
Extended by:
T::Sig
Defined in:
lib/rager/rerank/providers/cohere.rb

Instance Method Summary collapse

Instance Method Details

#rerank(input, options) ⇒ Object

Raises:



20
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
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
 
# File 'lib/rager/rerank/providers/cohere.rb', line 20

def rerank(input, options)
  api_key = options.api_key || ENV["COHERE_API_KEY"]
  raise Rager::Errors::CredentialsError.new("Cohere", env_var: ["COHERE_API_KEY"]) if api_key.nil?

  base_url = options.url || ENV["COHERE_URL"] || "https://api.cohere.com/v2"
  url = "#{base_url}/rerank"

  headers = {"Content-Type" => "application/json"}
  headers["Authorization"] = "Bearer #{api_key}" if api_key

  body = {
    model: options.model || "rerank-v3.5",
    query: input.query,
    documents: input.documents
  }.tap do |b|
    b[:top_n] = options.n if options.n
  end

  request = Rager::Http::Request.new(
    verb: Rager::Http::Verb::Post,
    url: url,
    headers: headers,
    body: body.to_json,
    timeout: options.timeout || Rager.config.timeout
  )

  response = Rager.config.http_adapter.make_request(request)
  response_body = T.cast(T.must(response.body), String)

  raise Rager::Errors::HttpError.new(Rager.config.http_adapter, request.url, response.status, body: response_body) if response.status != 200

  parsed_response = JSON.parse(response_body)
  results = parsed_response["results"]

  results.map do |r|
    score = r["relevance_score"].to_f
    content = input.documents[r["index"]]
    next if content.nil?

    Rager::Rerank::Result.new(
      score: score,
      content: content
    )
  end
end