Class: RerankerRuby::Cohere

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

Constant Summary collapse

API_URL = 
"https://api.cohere.com/v2/rerank"
DEFAULT_MODEL = 
"rerank-v3.5"

Instance Method Summary collapse

Constructor Details

#initialize(api_key:, model: DEFAULT_MODEL, **options) ⇒ Cohere

Returns a new instance of Cohere.



8
9
10
11
12
 
# File 'lib/reranker_ruby/cohere.rb', line 8

def initialize(api_key:, model: DEFAULT_MODEL, **options)
  super(**options)
  @api_key = api_key
  @model = model
end

Instance Method Details

#rerank(query, documents, top_k: 10, model: nil) ⇒ Object 



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
40
41
 
# File 'lib/reranker_ruby/cohere.rb', line 14

def rerank(query, documents, top_k: 10, model: nil)
  validate_inputs!(query, documents, top_k)
  instrument(query: query, document_count: documents.length, top_k: top_k) do
    with_cache(query, documents, top_k: top_k) do
      texts = extract_texts(documents)

      response = post(API_URL, {
        model: model || @model,
        query: query,
        documents: texts,
        top_n: top_k,
        return_documents: false
      }, headers: {
        "Authorization" => "Bearer #{@api_key}"
      })

      response["results"].map do |r|
        idx = r["index"]
        Result.new(
          text: texts[idx],
          score: r["relevance_score"],
          index: idx,
          metadata: extract_metadata(documents[idx])
        )
      end.sort
    end
  end
end