Class: RerankerRuby::Middleware

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

Overview

Pipeline middleware for integration with rag-ruby or any RAG pipeline. Accepts candidates from a retrieval step and reranks them.

Usage:

middleware = RerankerRuby::Middleware.new(
reranker: RerankerRuby::Cohere.new(api_key: "..."),
top_k: 5,
text_key: :content  # key to extract text from candidate hashes
)

reranked = middleware.call(query: "test", candidates: candidates)

Instance Method Summary collapse

Constructor Details

#initialize(reranker: nil, top_k: nil, text_key: :text) ⇒ Middleware

Returns a new instance of Middleware.



17
18
19
20
21
 
# File 'lib/reranker_ruby/middleware.rb', line 17

def initialize(reranker: nil, top_k: nil, text_key: :text)
  @reranker = reranker
  @top_k = top_k
  @text_key = text_key
end

Instance Method Details

#call(query:, candidates:, top_k: nil) ⇒ Object 



23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
 
# File 'lib/reranker_ruby/middleware.rb', line 23

def call(query:, candidates:, top_k: nil)
  reranker = @reranker || RerankerRuby.reranker
  top_k ||= @top_k || RerankerRuby.configuration.default_top_k

  documents = candidates.map do |candidate|
    case candidate
    when Hash
      text = candidate[@text_key] || candidate[@text_key.to_s]
      metadata = candidate.reject { |k, _| k == @text_key || k == @text_key.to_s }
      { text: text }.merge(metadata)
    when String
      candidate
    else
      # Duck-type: try to call the text_key method
      if candidate.respond_to?(@text_key)
        { text: candidate.send(@text_key) }
      else
        candidate.to_s
      end
    end
  end

  reranker.rerank(query, documents, top_k: top_k)
end