Module: RubyLLM::Protocols::Cohere::Rerank

Defined in:
lib/ruby_llm/protocols/cohere/rerank.rb

Overview

Rerank methods for the Cohere v2 API integration

Class Method Summary collapse

Class Method Details

.parse_rerank_response(response, model:, documents: []) ⇒ Object

Cohere returns positions and scores only, so the ranked text comes back from the documents that were sent.



25
26
27
28
29
30
31
32
33
34
35
36
 
# File 'lib/ruby_llm/protocols/cohere/rerank.rb', line 25

def parse_rerank_response(response, model:, documents: [])
  data = response.body
  billed = data.dig('meta', 'billed_units') || {}
  tokens = data.dig('meta', 'tokens') || {}

  RubyLLM::Rerank.new(
    results: parse_rerank_results(data, documents),
    model: model,
    raw: data,
    input_tokens: tokens['input_tokens'] || billed['input_tokens']
  )
end

.parse_rerank_results(data, documents) ⇒ Object 



38
39
40
41
42
43
44
45
46
47
48
49
50
51
 
# File 'lib/ruby_llm/protocols/cohere/rerank.rb', line 38

def parse_rerank_results(data, documents)
  Array(data['results']).map do |result|
    index = result['index']
    unless valid_rerank_index?(index, documents)
      raise Error, 'Cohere reranking returned an invalid document index'
    end

    RubyLLM::Rerank::Result.new(
      index: index,
      document: result.dig('document', 'text') || documents[index],
      score: result['relevance_score']
    )
  end
end

.render_rerank_payload(query, documents, model:, top_n: nil, provider_options: {}) ⇒ Object 



14
15
16
17
18
19
20
21
 
# File 'lib/ruby_llm/protocols/cohere/rerank.rb', line 14

def render_rerank_payload(query, documents, model:, top_n: nil, provider_options: {})
  {
    model: model,
    query: query,
    documents: documents,
    top_n: top_n
  }.compact.merge(provider_options)
end

.rerank_url ⇒ Object 



10
11
12
 
# File 'lib/ruby_llm/protocols/cohere/rerank.rb', line 10

def rerank_url
  'v2/rerank'
end

.valid_rerank_index?(index, documents) ⇒ Boolean

Returns:

  • (Boolean) 


53
54
55
 
# File 'lib/ruby_llm/protocols/cohere/rerank.rb', line 53

def valid_rerank_index?(index, documents)
  index.is_a?(Integer) && index.between?(0, documents.length - 1)
end