diff --git a/enterprise/app/services/captain/documents/hybrid_chunk_search_service.rb b/enterprise/app/services/captain/documents/hybrid_chunk_search_service.rb new file mode 100644 index 000000000..dfd0ed1ed --- /dev/null +++ b/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 diff --git a/enterprise/lib/captain/tools/faq_lookup_tool.rb b/enterprise/lib/captain/tools/faq_lookup_tool.rb index 93dd90259..06c78b48b 100644 --- a/enterprise/lib/captain/tools/faq_lookup_tool.rb +++ b/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 diff --git a/spec/enterprise/lib/captain/tools/faq_lookup_tool_spec.rb b/spec/enterprise/lib/captain/tools/faq_lookup_tool_spec.rb index ccae44ac2..2a64b0e61 100644 --- a/spec/enterprise/lib/captain/tools/faq_lookup_tool_spec.rb +++ b/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 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 new file mode 100644 index 000000000..bfffc3a07 --- /dev/null +++ b/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