diff --git a/enterprise/app/services/captain/documents/chunk_reranker_service.rb b/enterprise/app/services/captain/documents/chunk_reranker_service.rb new file mode 100644 index 000000000..3522c03a5 --- /dev/null +++ b/enterprise/app/services/captain/documents/chunk_reranker_service.rb @@ -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 diff --git a/enterprise/app/services/captain/documents/hybrid_chunk_search_service.rb b/enterprise/app/services/captain/documents/hybrid_chunk_search_service.rb index e508789ff..448063d8f 100644 --- a/enterprise/app/services/captain/documents/hybrid_chunk_search_service.rb +++ b/enterprise/app/services/captain/documents/hybrid_chunk_search_service.rb @@ -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 diff --git a/lib/llm_constants.rb b/lib/llm_constants.rb index 241f3ccba..3615eebb1 100644 --- a/lib/llm_constants.rb +++ b/lib/llm_constants.rb @@ -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-], diff --git a/spec/enterprise/services/captain/documents/chunk_reranker_service_spec.rb b/spec/enterprise/services/captain/documents/chunk_reranker_service_spec.rb new file mode 100644 index 000000000..26259284b --- /dev/null +++ b/spec/enterprise/services/captain/documents/chunk_reranker_service_spec.rb @@ -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 diff --git a/spec/enterprise/services/captain/documents/hybrid_chunk_search_service_spec.rb b/spec/enterprise/services/captain/documents/hybrid_chunk_search_service_spec.rb index da0f56c18..6302d4c24 100644 --- a/spec/enterprise/services/captain/documents/hybrid_chunk_search_service_spec.rb +++ b/spec/enterprise/services/captain/documents/hybrid_chunk_search_service_spec.rb @@ -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