feat(captain): improve contextual chunk retrieval quality

2026-02-20 10:16:46 +05:30
parent 8dc9789f12
commit 3df4dbf114
6 changed files with 275 additions and 15 deletions
@@ -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) }
@@ -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

@@ -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
@@ -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
@@ -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',
@@ -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