feat(captain): add chunking and context services

2026-02-19 14:27:50 +05:30
parent d3c91a29ca
commit 7c7200d98e
3 changed files with 256 additions and 0 deletions
@@ -0,0 +1,35 @@
+class Captain::Documents::ChunkEmbeddingService
+  def initialize(account_id:)
+    @account_id = account_id
+    @embedding_service = Captain::Llm::EmbeddingService.new(account_id: account_id)
+  end
+
+  def embed(content:, context: nil)
+    input = embedding_input(content: content, context: context)
+    return [] if input.blank?
+
+    @embedding_service.get_embedding(input)
+  end
+
+  def build_record_attributes(document:, chunk:, context: nil)
+    content = chunk.fetch(:content)
+    {
+      document_id: document.id,
+      assistant_id: document.assistant_id,
+      account_id: document.account_id,
+      position: chunk.fetch(:position),
+      content: content,
+      token_count: chunk[:token_count],
+      context: context,
+      embedding: embed(content: content, context: context),
+      created_at: Time.current,
+      updated_at: Time.current
+    }
+  end
+
+  private
+
+  def embedding_input(content:, context:)
+    [context, content].reject(&:blank?).join("\n\n")
+  end
+end
@@ -0,0 +1,144 @@
+class Captain::Documents::ChunkingService
+  DEFAULT_TARGET_TOKENS = 600
+  DEFAULT_MIN_TOKENS = 400
+  DEFAULT_MAX_TOKENS = 800
+  DEFAULT_OVERLAP_TOKENS = 120
+
+  def initialize(content, target_tokens: DEFAULT_TARGET_TOKENS, min_tokens: DEFAULT_MIN_TOKENS,
+                 max_tokens: DEFAULT_MAX_TOKENS, overlap_tokens: DEFAULT_OVERLAP_TOKENS)
+    @content = content.to_s
+    @target_tokens = target_tokens
+    @min_tokens = min_tokens
+    @max_tokens = max_tokens
+    @overlap_tokens = overlap_tokens
+  end
+
+  def chunk
+    return [] if @content.blank?
+
+    sections = split_into_sections(@content)
+    return [] if sections.empty?
+
+    base_chunks = build_chunks(sections)
+    apply_overlap(base_chunks).map.with_index do |chunk_content, index|
+      {
+        position: index,
+        content: chunk_content,
+        token_count: estimate_tokens(chunk_content)
+      }
+    end
+  end
+
+  private
+
+  Section = Struct.new(:content, :heading_path, keyword_init: true)
+
+  def split_into_sections(content)
+    heading_path = []
+
+    content
+      .split(/\n{2,}/)
+      .map(&:strip)
+      .reject(&:blank?)
+      .map do |section_content|
+      heading_path = update_heading_path_from_section(heading_path, section_content)
+      Section.new(content: section_content, heading_path: heading_path.dup)
+    end
+  end
+
+  def build_chunks(sections)
+    state = { chunks: [], current_chunk: +'', current_tokens: 0 }
+    sections.each { |section| process_section(section, state) }
+    state[:chunks] << state[:current_chunk] if state[:current_chunk].present?
+    state[:chunks]
+  end
+
+  def apply_overlap(chunks)
+    return chunks if chunks.size <= 1 || @overlap_tokens <= 0
+
+    overlapped = [chunks.first]
+    (1...chunks.size).each do |index|
+      previous_tail = tail_tokens(chunks[index - 1], @overlap_tokens)
+      next_chunk = [previous_tail, chunks[index]].reject(&:blank?).join("\n\n")
+      overlapped << next_chunk
+    end
+    overlapped
+  end
+
+  def with_heading_context(section)
+    return section.content if section.heading_path.empty?
+
+    heading_line = "Section: #{section.heading_path.join(' > ')}"
+    "#{heading_line}\n#{section.content}"
+  end
+
+  def heading?(line)
+    line.match?(/\A\#{1,6}\s+\S+/)
+  end
+
+  def appendable?(current_tokens, section_tokens)
+    return true if current_tokens + section_tokens <= @max_tokens
+    return false if current_tokens.zero?
+
+    current_tokens < @min_tokens
+  end
+
+  def append_to_current_chunk(current_chunk, section_content)
+    current_chunk << "\n\n" unless current_chunk.empty?
+    current_chunk << section_content
+    [current_chunk, estimate_tokens(current_chunk)]
+  end
+
+  def process_section(section, state)
+    section_content = with_heading_context(section)
+    section_tokens = estimate_tokens(section_content)
+    should_append = appendable?(state[:current_tokens], section_tokens)
+
+    state[:current_chunk], state[:current_tokens] = if should_append
+                                                      append_to_current_chunk(state[:current_chunk], section_content)
+                                                    else
+                                                      state[:chunks] << state[:current_chunk] if state[:current_chunk].present?
+                                                      [section_content.dup, section_tokens]
+                                                    end
+    flush_target_chunk(state)
+  end
+
+  def flush_target_chunk(state)
+    return unless state[:current_tokens] >= @target_tokens
+
+    state[:chunks] << state[:current_chunk]
+    state[:current_chunk] = +''
+    state[:current_tokens] = 0
+  end
+
+  def update_heading_path_from_section(path, section_content)
+    heading_line = section_content.lines.first.to_s.strip
+    return path unless heading?(heading_line)
+
+    update_heading_path(path, heading_line)
+  end
+
+  def update_heading_path(path, heading_line)
+    level = heading_line[/\A#+/].length
+    heading_text = heading_line.sub(/\A#+\s*/, '').strip
+
+    updated_path = path.dup
+    updated_path = updated_path.first(level - 1)
+    updated_path << heading_text
+    updated_path
+  end
+
+  def estimate_tokens(text)
+    return 0 if text.blank?
+
+    (text.split(/\s+/).length * 1.3).ceil
+  end
+
+  def tail_tokens(text, token_budget)
+    return '' if text.blank? || token_budget <= 0
+
+    words = text.split(/\s+/)
+    tail_word_count = (token_budget / 1.3).ceil
+    words.last(tail_word_count).join(' ')
+  end
+end
@@ -0,0 +1,77 @@
+class Captain::Documents::ContextGenerationService < Llm::BaseAiService
+  include Integrations::LlmInstrumentation
+
+  MAX_DOCUMENT_CHARACTERS = 20_000
+  MAX_CHUNK_CHARACTERS = 6_000
+  DEFAULT_MODEL = 'gpt-4.1-mini'.freeze
+
+  def initialize(document_content:, chunk_content:, account_id:, model: DEFAULT_MODEL)
+    super()
+    @document_content = document_content.to_s
+    @chunk_content = chunk_content.to_s
+    @account_id = account_id
+    @model = model
+  end
+
+  def generate
+    return '' if @chunk_content.blank?
+
+    response = instrument_llm_call(instrumentation_params) do
+      chat(model: @model, temperature: 0.1)
+        .with_instructions(system_prompt)
+        .ask(user_prompt)
+    end
+
+    response.content.to_s.strip
+  rescue RubyLLM::Error => e
+    Rails.logger.error "Context generation failed: #{e.message}"
+    ''
+  end
+
+  private
+
+  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.
+    PROMPT
+  end
+
+  def user_prompt
+    <<~PROMPT
+      <document>
+      #{truncated_document}
+      </document>
+
+      <chunk>
+      #{truncated_chunk}
+      </chunk>
+    PROMPT
+  end
+
+  def truncated_document
+    @document_content.first(MAX_DOCUMENT_CHARACTERS)
+  end
+
+  def truncated_chunk
+    @chunk_content.first(MAX_CHUNK_CHARACTERS)
+  end
+
+  def instrumentation_params
+    {
+      span_name: 'llm.captain.chunk_context',
+      model: @model,
+      temperature: 0.1,
+      feature_name: 'chunk_context_generation',
+      account_id: @account_id,
+      messages: [
+        { role: 'system', content: system_prompt },
+        { role: 'user', content: user_prompt }
+      ]
+    }
+  end
+end