akadmin/chatwoot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

feat(captain): add cohere reranker for chunk retrieval

Browse Source
This commit is contained in:
aakashb95
2026-02-20 17:15:20 +05:30
parent 25d30fbebf
commit 20e22897e6
5 changed files with 335 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/app/services/captain/documents/chunk_reranker_service.rb
+106
View File
@@ -0,0 +1,106 @@
class Captain::Documents::ChunkRerankerService
include Integrations::LlmInstrumentation
DEFAULT_MODEL = 'rerank-v4.0-pro'.freeze
DEFAULT_ENDPOINT = 'https://api.cohere.com/v2/rerank'.freeze
MAX_RERANK_TEXT_CHARACTERS = 1_200
def initialize(account_id:)
@account_id = account_id
@model = resolve_model
@api_key = InstallationConfig.find_by(name: 'CAPTAIN_COHERE_API_KEY')&.value.to_s
@endpoint = InstallationConfig.find_by(name: 'CAPTAIN_COHERE_API_BASE')&.value.to_s.presence || DEFAULT_ENDPOINT
end
def rerank(query:, candidates:, limit:)
return candidates.first(limit) if candidates.blank?
return candidates.first(limit) if @api_key.blank?
reranked_indices = fetch_reranked_indices(query, candidates)
return candidates.first(limit) if reranked_indices.blank?
reranked_chunks = reranked_indices.filter_map { |index| candidates[index] }
remaining_chunks = candidates.reject { |chunk| reranked_chunks.include?(chunk) }
(reranked_chunks + remaining_chunks).first(limit)
rescue StandardError => e
Rails.logger.warn "Chunk reranker failed: #{e.message}"
candidates.first(limit)
end
private
def reranker_text(chunk)
[chunk.context, chunk.content]
.compact
.join("\n\n")
.squish
.first(MAX_RERANK_TEXT_CHARACTERS)
end
def instrumentation_params(query, documents)
{
span_name: 'llm.captain.chunk_rerank',
model: @model,
feature_name: 'chunk_rerank',
account_id: @account_id,
messages: [
{ role: 'user', content: { query: query, document_count: documents.size }.to_json }
]
}
end
def resolve_model
InstallationConfig.find_by(name: 'CAPTAIN_CHUNK_RERANK_MODEL')&.value.presence || DEFAULT_MODEL
end
def fetch_reranked_indices(query, candidates)
documents = candidates.map { |chunk| reranker_text(chunk) }
response_payload = instrument_llm_call(instrumentation_params(query, documents)) do
cohere_rerank(query, documents)
end
parse_reranked_indices(response_payload, candidates.size)
end
def cohere_rerank(query, documents)
response = HTTParty.post(
@endpoint,
body: {
model: @model,
query: query,
documents: documents,
top_n: documents.size
}.to_json,
headers: cohere_headers
)
raise "Cohere rerank request failed with status #{response.code}" unless response.success?
parsed = response.parsed_response
raise 'Cohere rerank response is invalid' unless parsed.is_a?(Hash)
parsed
end
def parse_reranked_indices(payload, candidate_count)
results = payload['results']
return [] unless results.is_a?(Array)
results
.filter_map do |result|
index = Integer(result['index'], exception: false)
next nil if index.blank?
next nil if index.negative? || index >= candidate_count
index
end
.uniq
end
def cohere_headers
{
'Authorization' => "Bearer #{@api_key}",
'Content-Type' => 'application/json'
}
end
end
enterprise/app/services/captain/documents/hybrid_chunk_search_service.rb
+27 -3
View File
@@ -1,8 +1,10 @@
class Captain::Documents::HybridChunkSearchService
DEFAULT_VECTOR_LIMIT = 10
DEFAULT_BM25_LIMIT = 10
DEFAULT_VECTOR_LIMIT = 20
DEFAULT_BM25_LIMIT = 20
DEFAULT_BM25_CANDIDATE_LIMIT = 200
DEFAULT_RRF_LIMIT = 20
DEFAULT_RESULT_LIMIT = 5
DEFAULT_RERANKER_ENABLED = true
RRF_K = 60
BM25_K1 = 1.2
BM25_B = 0.75
@@ -19,6 +21,7 @@ class Captain::Documents::HybridChunkSearchService
def initialize(assistant:)
@assistant = assistant
@embedding_service = Captain::Llm::EmbeddingService.new(account_id: assistant.account_id)
@reranker = Captain::Documents::ChunkRerankerService.new(account_id: assistant.account_id)
end
def search(query, limit: DEFAULT_RESULT_LIMIT)
@@ -26,7 +29,8 @@ class Captain::Documents::HybridChunkSearchService
vector_results = vector_search(query)
bm25_results = bm25_search(query)
rank_results(vector_results, bm25_results, limit)
rrf_results = rank_results(vector_results, bm25_results, DEFAULT_RRF_LIMIT)
rerank_results(query, rrf_results, limit)
end
private
@@ -80,6 +84,19 @@ class Captain::Documents::HybridChunkSearchService
.filter_map { |(id, _score)| records_by_id[id] }
end
def rerank_results(query, candidates, limit)
return candidates.first(limit) unless reranking_enabled?
return candidates.first(limit) if candidates.blank?
reranked_results = @reranker.rerank(query: query, candidates: candidates, limit: limit)
return candidates.first(limit) if reranked_results.blank?
reranked_results
rescue StandardError => e
Rails.logger.warn "Hybrid chunk rerank failed: #{e.message}"
candidates.first(limit)
end
def bm25_candidates(query_terms)
tsquery = build_or_tsquery(query_terms)
return [] if tsquery.blank?
@@ -185,4 +202,11 @@ class Captain::Documents::HybridChunkSearchService
def sanitize_tsquery_term(term)
term.to_s.gsub(/[^a-z0-9]/, '')
end
def reranking_enabled?
value = InstallationConfig.find_by(name: 'CAPTAIN_CHUNK_RERANKING_ENABLED')&.value
return DEFAULT_RERANKER_ENABLED if value.blank?
ActiveModel::Type::Boolean.new.cast(value)
end
end
lib/llm_constants.rb
+1
View File
@@ -9,6 +9,7 @@ module LlmConstants
PROVIDER_PREFIXES = {
'openai' => %w[gpt- o1 o3 o4 text-embedding- whisper- tts-],
'cohere' => %w[rerank- command- embed-],
'anthropic' => %w[claude-],
'google' => %w[gemini-],
'mistral' => %w[mistral- codestral-],
spec/enterprise/services/captain/documents/chunk_reranker_service_spec.rb
+114
View File
@@ -0,0 +1,114 @@
require 'rails_helper'
RSpec.describe Captain::Documents::ChunkRerankerService do
let(:assistant) { create(:captain_assistant) }
let(:account) { assistant.account }
let(:service) { described_class.new(account_id: account.id) }
before do
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_COHERE_API_KEY').update!(value: 'cohere-test-key')
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_CHUNK_RERANK_MODEL').update!(value: 'rerank-v4.0-pro')
end
describe '#rerank' do
it 'returns chunks in reranked order when the model returns ids' do
document = create(:captain_document, account: account, assistant: assistant, chunking_status: :ready, status: :available)
first_chunk = create(
:captain_document_chunk,
document: document,
account: account,
assistant: assistant,
position: 0,
content: 'How to cancel a subscription'
)
second_chunk = create(
:captain_document_chunk,
document: document,
account: account,
assistant: assistant,
position: 1,
content: 'How to delete an account permanently'
)
response_payload = {
'results' => [
{ 'index' => 1, 'relevance_score' => 0.93 },
{ 'index' => 0, 'relevance_score' => 0.67 }
]
}
cohere_response = instance_double(
HTTParty::Response,
success?: true,
code: 200,
parsed_response: response_payload
)
allow(HTTParty).to receive(:post).and_return(cohere_response)
results = service.rerank(query: 'delete account', candidates: [first_chunk, second_chunk], limit: 2)
expect(results.map(&:id)).to eq([second_chunk.id, first_chunk.id])
end
it 'falls back to original order when response payload is invalid' do
document = create(:captain_document, account: account, assistant: assistant, chunking_status: :ready, status: :available)
first_chunk = create(
:captain_document_chunk,
document: document,
account: account,
assistant: assistant,
position: 0,
content: 'How to cancel a subscription'
)
second_chunk = create(
:captain_document_chunk,
document: document,
account: account,
assistant: assistant,
position: 1,
content: 'How to delete an account permanently'
)
cohere_response = instance_double(
HTTParty::Response,
success?: true,
code: 200,
parsed_response: { 'results' => nil }
)
allow(HTTParty).to receive(:post).and_return(cohere_response)
results = service.rerank(query: 'delete account', candidates: [first_chunk, second_chunk], limit: 2)
expect(results.map(&:id)).to eq([first_chunk.id, second_chunk.id])
end
it 'falls back to original order when API key is missing' do
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_COHERE_API_KEY').update!(value: '')
document = create(:captain_document, account: account, assistant: assistant, chunking_status: :ready, status: :available)
first_chunk = create(
:captain_document_chunk,
document: document,
account: account,
assistant: assistant,
position: 0,
content: 'How to cancel a subscription'
)
second_chunk = create(
:captain_document_chunk,
document: document,
account: account,
assistant: assistant,
position: 1,
content: 'How to delete an account permanently'
)
results = described_class.new(account_id: account.id).rerank(
query: 'delete account',
candidates: [first_chunk, second_chunk],
limit: 2
)
expect(results.map(&:id)).to eq([first_chunk.id, second_chunk.id])
end
end
end
spec/enterprise/services/captain/documents/hybrid_chunk_search_service_spec.rb
+87
View File
@@ -7,6 +7,7 @@ RSpec.describe Captain::Documents::HybridChunkSearchService do
before do
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_OPEN_AI_API_KEY').update!(value: 'test-key')
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_CHUNK_RERANKING_ENABLED').update!(value: 'false')
end
describe '#search' do
@@ -91,5 +92,91 @@ RSpec.describe Captain::Documents::HybridChunkSearchService do
expect(results.first.id).to eq(stronger_chunk.id)
expect(results.map(&:id)).to include(weaker_chunk.id)
end
it 'reranks RRF candidates with the configured reranker' do
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_CHUNK_RERANKING_ENABLED').update!(value: 'true')
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_CHUNK_RERANK_MODEL').update!(value: 'rerank-v4.0-pro')
ready_document = create(
:captain_document,
account: account,
assistant: assistant,
status: :available,
chunking_status: :ready
)
first_chunk = create(
:captain_document_chunk,
document: ready_document,
account: account,
assistant: assistant,
position: 0,
content: 'Cancel subscription from billing settings on iOS and Android.'
)
second_chunk = create(
:captain_document_chunk,
document: ready_document,
account: account,
assistant: assistant,
position: 1,
content: 'Delete your account permanently from profile settings and confirm with password.'
)
embedding_service = instance_double(Captain::Llm::EmbeddingService, get_embedding: [])
allow(Captain::Llm::EmbeddingService).to receive(:new).and_return(embedding_service)
reranker_service = instance_double(
Captain::Documents::ChunkRerankerService,
rerank: [second_chunk, first_chunk]
)
allow(Captain::Documents::ChunkRerankerService).to receive(:new).and_return(reranker_service)
results = described_class.new(assistant: assistant).search('How do I cancel or delete my account?', limit: 2)
expect(results.map(&:id)).to eq([second_chunk.id, first_chunk.id])
end
it 'falls back to RRF order when reranking fails' do
ready_document = create(
:captain_document,
account: account,
assistant: assistant,
status: :available,
chunking_status: :ready
)
create(
:captain_document_chunk,
document: ready_document,
account: account,
assistant: assistant,
position: 0,
content: 'Incognito mode exists for profile visibility.',
context: 'Privacy and safety settings.'
)
create(
:captain_document_chunk,
document: ready_document,
account: account,
assistant: assistant,
position: 1,
content: 'Incognito mode allows hidden browsing. Incognito mode keeps your profile hidden.',
context: 'Incognito mode details and hidden profile behavior.'
)
embedding_service = instance_double(Captain::Llm::EmbeddingService, get_embedding: [])
allow(Captain::Llm::EmbeddingService).to receive(:new).and_return(embedding_service)
fallback_results = described_class.new(assistant: assistant).search('how does incognito mode work', limit: 2)
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_CHUNK_RERANKING_ENABLED').update!(value: 'true')
reranker_service = instance_double(Captain::Documents::ChunkRerankerService)
allow(reranker_service).to receive(:rerank).and_raise(StandardError, 'reranker unavailable')
allow(Captain::Documents::ChunkRerankerService).to receive(:new).and_return(reranker_service)
reranker_failed_results = described_class.new(assistant: assistant).search('how does incognito mode work', limit: 2)
expect(reranker_failed_results.map(&:id)).to eq(fallback_results.map(&:id))
end
end
end