21 def rerank(query, documents, options)
22 api_key = options.api_key || ENV["COHERE_API_KEY"]
23 raise Rager::Errors::CredentialsError.new("Cohere", env_var: ["COHERE_API_KEY"]) if api_key.nil?
24
25 base_url = options.url || ENV["COHERE_URL"] || "https://api.cohere.com/v2"
26
27 body = {
28 model: options.model || "rerank-v3.5",
29 query: query,
30 documents: documents
31 }.tap do |b|
32 b[:top_n] = options.n if options.n
33 end
34
35 headers = {"Content-Type" => "application/json"}
36 headers["Authorization"] = "Bearer #{api_key}" if api_key
37
38 request = Rager::Http::Request.new(
39 verb: Rager::Http::Verb::Post,
40 url: "#{base_url}/rerank",
41 headers: headers,
42 body: body.to_json,
43 timeout: options.timeout || Rager.config.timeout
44 )
45
46 response = Rager.config.http_adapter.make_request(request)
47 response_body = T.cast(T.must(response.body), String)
48
49 raise Rager::Errors::HttpError.new(Rager.config.http_adapter, request.url, response.status, body: response_body) if response.status != 200
50
51 parsed_response = JSON.parse(response_body)
52 results = parsed_response["results"] || []
53
54 scores = results.map { |r| r["relevance_score"].to_f }
55 reranked_documents = results.map { |r| documents[r["index"]] }
56
57 Rager::Rerank::Output.new(scores: scores, documents: reranked_documents)
58 end