feat(captain): add chunk retrieval mode for faq lookup

enterprise/app/services/captain/documents/hybrid_chunk_search_service.rb

@@ -0,0 +1,70 @@
+class Captain::Documents::HybridChunkSearchService
+  DEFAULT_VECTOR_LIMIT = 10
+  DEFAULT_BM25_LIMIT = 10
+  DEFAULT_RESULT_LIMIT = 5
+  RRF_K = 60
+
+  def initialize(assistant:)
+    @assistant = assistant
+    @embedding_service = Captain::Llm::EmbeddingService.new(account_id: assistant.account_id)
+  end
+
+  def search(query, limit: DEFAULT_RESULT_LIMIT)
+    return [] if query.blank?
+
+    vector_results = vector_search(query)
+    bm25_results = bm25_search(query)
+    rank_results(vector_results, bm25_results, limit)
+  end
+
+  private
+
+  def base_scope
+    Captain::DocumentChunk
+      .where(account_id: @assistant.account_id, assistant_id: @assistant.id)
+      .joins(:document)
+      .merge(Captain::Document.chunking_status_ready)
+      .includes(:document)
+  end
+
+  def vector_search(query)
+    embedding = @embedding_service.get_embedding(query)
+    return [] if embedding.blank?
+
+    base_scope
+      .all
+      .nearest_neighbors(:embedding, embedding, distance: 'cosine')
+      .limit(DEFAULT_VECTOR_LIMIT)
+      .to_a
+  end
+
+  def bm25_search(query)
+    scope = base_scope
+    quoted_query = ActiveRecord::Base.connection.quote(query)
+    text_expression = "to_tsvector('english', coalesce(captain_document_chunks.context, '') || ' ' || captain_document_chunks.content)"
+
+    scope
+      .where(Arel.sql("#{text_expression} @@ plainto_tsquery('english', #{quoted_query})"))
+      .order(Arel.sql("ts_rank_cd(#{text_expression}, plainto_tsquery('english', #{quoted_query})) DESC"))
+      .limit(DEFAULT_BM25_LIMIT)
+      .to_a
+  end
+
+  def rank_results(vector_results, bm25_results, limit)
+    score_by_id = Hash.new(0.0)
+
+    vector_results.each_with_index do |chunk, index|
+      score_by_id[chunk.id] += 1.0 / (RRF_K + index + 1)
+    end
+
+    bm25_results.each_with_index do |chunk, index|
+      score_by_id[chunk.id] += 1.0 / (RRF_K + index + 1)
+    end
+
+    records_by_id = (vector_results + bm25_results).index_by(&:id)
+    score_by_id
+      .sort_by { |(_id, score)| -score }
+      .first(limit)
+      .filter_map { |(id, _score)| records_by_id[id] }
+  end
+end

enterprise/lib/captain/tools/faq_lookup_tool.rb

@@ -5,20 +5,73 @@ class Captain::Tools::FaqLookupTool < Captain::Tools::BasePublicTool
   def perform(_tool_context, query:)
     log_tool_usage('searching', { query: query })
 
-    # Use existing vector search on approved responses
-    responses = @assistant.responses.approved.search(query).to_a
+    faq_results, chunk_results = search_knowledge(query)
+    total_results = faq_results.size + chunk_results.size
 
-    if responses.empty?
+    if total_results.zero?
       log_tool_usage('no_results', { query: query })
       "No relevant FAQs found for: #{query}"
     else
-      log_tool_usage('found_results', { query: query, count: responses.size })
-      format_responses(responses)
+      log_tool_usage('found_results', { query: query, count: total_results })
+      "#{format_chunk_results(chunk_results)}#{format_responses(faq_results)}"
     end
   end
 
   private
 
+  def search_knowledge(query)
+    if chunk_retrieval_mode?
+      [
+        search_non_document_faqs(query),
+        Captain::Documents::HybridChunkSearchService.new(assistant: @assistant).search(query)
+      ]
+    else
+      [@assistant.responses.approved.search(query).to_a, []]
+    end
+  end
+
+  def search_non_document_faqs(query)
+    @assistant.responses
+              .approved
+              .where.not(documentable_type: 'Captain::Document')
+              .search(query)
+              .to_a
+  end
+
+  def chunk_retrieval_mode?
+    return false unless chunk_builder_enabled?
+
+    value = @assistant.config&.fetch('feature_document_faq_generation', true)
+    !ActiveModel::Type::Boolean.new.cast(value)
+  end
+
+  def chunk_builder_enabled?
+    value = InstallationConfig.find_by(name: 'CAPTAIN_DOCUMENT_CHUNKING_ENABLED')&.value
+    ActiveModel::Type::Boolean.new.cast(value)
+  end
+
+  def format_chunk_results(chunks)
+    chunks.map { |chunk| format_chunk(chunk) }.join
+  end
+
+  def format_chunk(chunk)
+    document = chunk.document
+    source_link = document.external_link if should_show_document_source?(document)
+    title = document.name.presence || document.external_link
+
+    formatted = "
+        Article: #{title}
+        Context: #{chunk.context}
+        Content: #{chunk.content}
+        "
+    if source_link.present?
+      formatted += "
+        Source: #{source_link}
+        "
+    end
+    formatted
+  end
+
   def format_responses(responses)
     responses.map { |response| format_response(response) }.join
   end
@@ -45,4 +98,11 @@ class Captain::Tools::FaqLookupTool < Captain::Tools::BasePublicTool
     external_link = response.documentable.external_link
     !external_link.start_with?('PDF:')
   end
+
+  def should_show_document_source?(document)
+    return false if document.blank?
+    return false if document.external_link.blank?
+
+    !document.external_link.start_with?('PDF:')
+  end
 end

spec/enterprise/lib/captain/tools/faq_lookup_tool_spec.rb

@@ -110,6 +110,48 @@ RSpec.describe Captain::Tools::FaqLookupTool, type: :model do
         expect(result).to eq('No relevant FAQs found for: ')
       end
     end
+
+    context 'when chunk retrieval mode is enabled for assistant' do
+      let(:document) do
+        create(
+          :captain_document,
+          assistant: assistant,
+          account: account,
+          chunking_status: :ready,
+          status: :available,
+          external_link: 'https://help.example.com/pricing',
+          name: 'Pricing'
+        )
+      end
+      let(:chunk) do
+        create(
+          :captain_document_chunk,
+          document: document,
+          assistant: assistant,
+          account: account,
+          content: 'Business plan starts at $19.',
+          context: 'Pricing page details'
+        )
+      end
+      let(:chunk_search_service) { instance_double(Captain::Documents::HybridChunkSearchService) }
+
+      before do
+        assistant.update!(config: (assistant.config || {}).merge('feature_document_faq_generation' => false))
+        InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_DOCUMENT_CHUNKING_ENABLED').update!(value: 'true')
+        allow(Captain::Documents::HybridChunkSearchService).to receive(:new).with(assistant: assistant).and_return(chunk_search_service)
+        allow(chunk_search_service).to receive(:search).with('pricing').and_return([chunk])
+        allow(Captain::AssistantResponse).to receive(:nearest_neighbors).and_return(Captain::AssistantResponse.none)
+      end
+
+      it 'returns chunk content for knowledge lookup' do
+        result = tool.perform(tool_context, query: 'pricing')
+
+        expect(result).to include('Article: Pricing')
+        expect(result).to include('Context: Pricing page details')
+        expect(result).to include('Content: Business plan starts at $19.')
+        expect(result).to include('Source: https://help.example.com/pricing')
+      end
+    end
   end
 
   describe '#active?' do

spec/enterprise/services/captain/documents/hybrid_chunk_search_service_spec.rb

@@ -0,0 +1,58 @@
+require 'rails_helper'
+
+RSpec.describe Captain::Documents::HybridChunkSearchService do
+  let(:assistant) { create(:captain_assistant) }
+  let(:account) { assistant.account }
+  let(:service) { described_class.new(assistant: assistant) }
+
+  before do
+    InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_OPEN_AI_API_KEY').update!(value: 'test-key')
+  end
+
+  describe '#search' do
+    it 'returns no results for blank query' do
+      expect(service.search('')).to eq([])
+    end
+
+    it 'returns chunk results from ready documents scoped to the assistant' do
+      ready_document = create(
+        :captain_document,
+        account: account,
+        assistant: assistant,
+        status: :available,
+        chunking_status: :ready
+      )
+      skipped_document = create(
+        :captain_document,
+        account: account,
+        assistant: assistant,
+        status: :available,
+        chunking_status: :pending
+      )
+
+      create(
+        :captain_document_chunk,
+        document: ready_document,
+        account: account,
+        assistant: assistant,
+        content: 'How to reset password for admins'
+      )
+      create(
+        :captain_document_chunk,
+        document: skipped_document,
+        account: account,
+        assistant: assistant,
+        content: 'How to reset password for members'
+      )
+
+      embedding_service = instance_double(Captain::Llm::EmbeddingService, get_embedding: [])
+      allow(Captain::Llm::EmbeddingService).to receive(:new).and_return(embedding_service)
+
+      results = service.search('reset password')
+
+      expect(results.size).to eq(1)
+      expect(results.first.document_id).to eq(ready_document.id)
+      expect(results.first.content).to include('reset password')
+    end
+  end
+end