feat(captain): add cohere reranker for chunk retrieval

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

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