From 3df4dbf1145bf27d100ca083bf50abc3acd9c216 Mon Sep 17 00:00:00 2001 From: aakashb95 Date: Fri, 20 Feb 2026 10:16:46 +0530 Subject: [PATCH] feat(captain): improve contextual chunk retrieval quality --- .../captain/documents/chunking_service.rb | 67 ++++++++- .../documents/context_generation_service.rb | 14 +- .../documents/hybrid_chunk_search_service.rb | 134 ++++++++++++++++-- .../documents/chunking_service_spec.rb | 25 ++++ .../context_generation_service_spec.rb | 13 ++ .../hybrid_chunk_search_service_spec.rb | 37 +++++ 6 files changed, 275 insertions(+), 15 deletions(-) diff --git a/enterprise/app/services/captain/documents/chunking_service.rb b/enterprise/app/services/captain/documents/chunking_service.rb index dc6f827f0..aaee36a20 100644 --- a/enterprise/app/services/captain/documents/chunking_service.rb +++ b/enterprise/app/services/captain/documents/chunking_service.rb @@ -3,6 +3,19 @@ class Captain::Documents::ChunkingService DEFAULT_MIN_TOKENS = 400 DEFAULT_MAX_TOKENS = 800 DEFAULT_OVERLAP_TOKENS = 120 + BOILERPLATE_SECTION_PATTERNS = [ + /skip to main content/i, + /table of contents/i, + /related articles/i, + /recommended articles/i, + /need more help/i, + /contact support/i, + /cookie/i, + /privacy policy/i, + /terms of service/i, + /all rights reserved/i, + /back to top/i + ].freeze def initialize(content, target_tokens: DEFAULT_TARGET_TOKENS, min_tokens: DEFAULT_MIN_TOKENS, max_tokens: DEFAULT_MAX_TOKENS, overlap_tokens: DEFAULT_OVERLAP_TOKENS) @@ -34,9 +47,10 @@ class Captain::Documents::ChunkingService Section = Struct.new(:content, :heading_path, keyword_init: true) def split_into_sections(content) + cleaned_content = remove_boilerplate_sections(content) heading_path = [] - content + cleaned_content .split(/\n{2,}/) .map(&:strip) .reject(&:blank?) @@ -46,6 +60,57 @@ class Captain::Documents::ChunkingService end end + def remove_boilerplate_sections(content) + content + .split(/\n{2,}/) + .map(&:strip) + .reject(&:blank?) + .reject { |section| boilerplate_section?(section) } + .join("\n\n") + end + + def boilerplate_section?(section) + normalized = section.downcase.strip + return true if BOILERPLATE_SECTION_PATTERNS.any? { |pattern| normalized.match?(pattern) } + + link_heavy_navigation_section?(section) + end + + def link_heavy_navigation_section?(section) + lines = non_blank_lines(section) + return false unless navigation_candidate?(lines) + + markdown_links = markdown_link_count(section) + linked_lines = linked_line_count(lines) + return false unless dense_link_cluster?(markdown_links, linked_lines) + + short_section?(section) + end + + def non_blank_lines(section) + section.lines.map(&:strip).reject(&:blank?) + end + + def navigation_candidate?(lines) + lines.size >= 3 + end + + def markdown_link_count(section) + section.scan(/\[[^\]]+\]\([^)]+\)/).size + end + + def linked_line_count(lines) + lines.count { |line| line.match?(/\[[^\]]+\]\([^)]+\)/) || line.start_with?('* [', '- [') } + end + + def dense_link_cluster?(markdown_links, linked_lines) + markdown_links >= 2 && linked_lines >= 3 + end + + def short_section?(section) + section.scan(/\b[\w']+\b/).size <= 180 + end + def build_chunks(sections) state = { chunks: [], current_chunk: +'', current_tokens: 0 } sections.each { |section| process_section(section, state) } diff --git a/enterprise/app/services/captain/documents/context_generation_service.rb b/enterprise/app/services/captain/documents/context_generation_service.rb index fee5f1580..1a5307fc5 100644 --- a/enterprise/app/services/captain/documents/context_generation_service.rb +++ b/enterprise/app/services/captain/documents/context_generation_service.rb @@ -32,12 +32,14 @@ class Captain::Documents::ContextGenerationService < Llm::BaseAiService def system_prompt <<~PROMPT - You generate retrieval context for document chunks. - Return 2 to 3 sentences that explain: - - which page/section this chunk belongs to - - what the chunk is mainly about - - key entities, plans, features, or limits mentioned - Keep it factual and concise. + You are writing retrieval context for a knowledge-base chunk. + Output 2 concise sentences that make this chunk easier to find for user questions. + Focus on: + - user intents this chunk can answer + - product terms and alternate phrasings users may search for + - key actions, settings, limits, or troubleshooting signals + Keep it factual. Do not mention "this chunk" or "this section". + Do not add information that is not supported by the document or chunk. PROMPT 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 dfd0ed1ed..e508789ff 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,20 @@ class Captain::Documents::HybridChunkSearchService DEFAULT_VECTOR_LIMIT = 10 DEFAULT_BM25_LIMIT = 10 + DEFAULT_BM25_CANDIDATE_LIMIT = 200 DEFAULT_RESULT_LIMIT = 5 RRF_K = 60 + BM25_K1 = 1.2 + BM25_B = 0.75 + ENGLISH_STOPWORDS = Set.new( + %w[ + a an and are as at be by for from has have how i in is it of on or that the this to was what when + where who why with your our their they we you do does did can could should would may might will + about above after again against all am any because before being below between both but down during + each few further here into more most other over own same some such than then there these those + through under until up very while + ] + ).freeze def initialize(assistant:) @assistant = assistant @@ -39,15 +51,15 @@ class Captain::Documents::HybridChunkSearchService 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)" + query_terms = bm25_terms(query).uniq + return [] if query_terms.blank? - 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 + candidates = bm25_candidates(query_terms) + return [] if candidates.blank? + + score_bm25_candidates(candidates, query_terms) + .first(DEFAULT_BM25_LIMIT) + .map(&:first) end def rank_results(vector_results, bm25_results, limit) @@ -67,4 +79,110 @@ class Captain::Documents::HybridChunkSearchService .first(limit) .filter_map { |(id, _score)| records_by_id[id] } end + + def bm25_candidates(query_terms) + tsquery = build_or_tsquery(query_terms) + return [] if tsquery.blank? + + quoted_tsquery = ActiveRecord::Base.connection.quote(tsquery) + vector_expression = "to_tsvector('english', coalesce(captain_document_chunks.context, '') || ' ' || captain_document_chunks.content)" + + base_scope + .where(Arel.sql("#{vector_expression} @@ to_tsquery('english', #{quoted_tsquery})")) + .limit(DEFAULT_BM25_CANDIDATE_LIMIT) + .to_a + end + + def score_bm25_candidates(chunks, query_terms) + documents = chunks.filter_map { |chunk| bm25_document_payload(chunk, query_terms) } + return [] if documents.blank? + + scoring_context = bm25_scoring_context(documents, query_terms) + scores = documents.map do |document| + [document[:chunk], bm25_document_score(document, scoring_context)] + end + scores.sort_by { |(_chunk, score)| -score } + end + + def bm25_document_payload(chunk, query_terms) + terms = bm25_terms([chunk.context, chunk.content].join(' ')) + return nil if terms.blank? + + term_frequencies = terms.tally.slice(*query_terms) + return nil if term_frequencies.blank? + + { + chunk: chunk, + doc_length: terms.length, + term_freqs: term_frequencies + } + end + + def bm25_document_frequency(documents, query_terms) + query_terms.index_with do |term| + documents.count { |document| document[:term_freqs].key?(term) } + end + end + + def average_document_length(documents) + total_tokens = documents.sum { |document| document[:doc_length] } + return 1.0 if total_tokens.zero? + + total_tokens.to_f / documents.size + end + + def bm25_scoring_context(documents, query_terms) + { + avg_doc_length: average_document_length(documents), + doc_frequency: bm25_document_frequency(documents, query_terms), + total_documents: documents.size, + query_terms: query_terms + } + end + + def bm25_document_score(document, scoring_context) + term_frequencies = document[:term_freqs] + scoring_context[:query_terms].sum do |term| + bm25_term_score(term, term_frequencies[term].to_i, document[:doc_length], scoring_context) + end + end + + def bm25_term_score(term, term_frequency, doc_length, scoring_context) + return 0.0 if term_frequency.zero? + + document_frequency = scoring_context[:doc_frequency][term].to_i + return 0.0 if document_frequency.zero? + + idf = bm25_inverse_document_frequency(document_frequency, scoring_context[:total_documents]) + denominator = bm25_denominator(term_frequency, doc_length, scoring_context[:avg_doc_length]) + idf * ((term_frequency * (BM25_K1 + 1.0)) / denominator) + end + + def bm25_inverse_document_frequency(document_frequency, total_documents) + Math.log(1.0 + ((total_documents - document_frequency + 0.5) / (document_frequency + 0.5))) + end + + def bm25_denominator(term_frequency, doc_length, avg_doc_length) + term_frequency + (BM25_K1 * (1.0 - BM25_B + (BM25_B * doc_length / avg_doc_length))) + end + + def bm25_terms(text) + text + .to_s + .downcase + .scan(/[a-z0-9]+/) + .reject { |token| token.length < 2 || ENGLISH_STOPWORDS.include?(token) } + end + + def build_or_tsquery(query_terms) + query_terms + .map { |term| sanitize_tsquery_term(term) } + .reject(&:blank?) + .map { |term| "#{term}:*" } + .join(' | ') + end + + def sanitize_tsquery_term(term) + term.to_s.gsub(/[^a-z0-9]/, '') + end end diff --git a/spec/enterprise/services/captain/documents/chunking_service_spec.rb b/spec/enterprise/services/captain/documents/chunking_service_spec.rb index 72f2f61e3..5ddba777e 100644 --- a/spec/enterprise/services/captain/documents/chunking_service_spec.rb +++ b/spec/enterprise/services/captain/documents/chunking_service_spec.rb @@ -53,5 +53,30 @@ RSpec.describe Captain::Documents::ChunkingService do expect(result.size).to eq(2) expect(result.last[:content]).to include('nine ten') end + + it 'removes obvious boilerplate navigation sections before chunking' do + content = <<~TEXT + # Account deletion + You can delete your account from Settings. + + ## Related articles + - [How to block someone](https://example.com/block) + - [How to report someone](https://example.com/report) + - [How to update profile](https://example.com/profile) + TEXT + + result = described_class.new( + content, + target_tokens: 30, + min_tokens: 10, + max_tokens: 50, + overlap_tokens: 0 + ).chunk + + combined_content = result.map { |chunk| chunk[:content] }.join("\n") + expect(combined_content).to include('You can delete your account from Settings') + expect(combined_content).not_to include('How to block someone') + expect(combined_content).not_to include('Related articles') + end end end diff --git a/spec/enterprise/services/captain/documents/context_generation_service_spec.rb b/spec/enterprise/services/captain/documents/context_generation_service_spec.rb index 248bc19cc..e8c7e983a 100644 --- a/spec/enterprise/services/captain/documents/context_generation_service_spec.rb +++ b/spec/enterprise/services/captain/documents/context_generation_service_spec.rb @@ -28,6 +28,19 @@ RSpec.describe Captain::Documents::ContextGenerationService do expect(result).to eq('Pricing page context.') end + it 'uses retrieval-oriented context instructions' do + service = described_class.new( + document_content: 'Document text', + chunk_content: 'Chunk text', + account_id: 1 + ) + + service.generate + + expect(chat).to have_received(:with_instructions).with(include('user intents this chunk can answer')) + expect(chat).to have_received(:with_instructions).with(include('make this chunk easier to find')) + end + it 'uses explicit model when provided' do service = described_class.new( document_content: 'Doc text', 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 bfffc3a07..da0f56c18 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 @@ -54,5 +54,42 @@ RSpec.describe Captain::Documents::HybridChunkSearchService do expect(results.first.document_id).to eq(ready_document.id) expect(results.first.content).to include('reset password') end + + it 'uses BM25 scoring to prioritize stronger lexical matches' do + ready_document = create( + :captain_document, + account: account, + assistant: assistant, + status: :available, + chunking_status: :ready + ) + + weaker_chunk = 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.' + ) + stronger_chunk = 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) + + results = service.search('how does incognito mode work') + + expect(results.first.id).to eq(stronger_chunk.id) + expect(results.map(&:id)).to include(weaker_chunk.id) + end end end