test(captain): add behavior specs for chunk services

db/migrate/20260219143000_add_captain_document_chunks_and_chunking_progress.rb

@@ -0,0 +1,65 @@
+class AddCaptainDocumentChunksAndChunkingProgress < ActiveRecord::Migration[7.0]
+  def up
+    create_document_chunks_table
+    add_chunking_progress_columns
+  end
+
+  def down
+    remove_chunking_progress_columns
+    drop_table :captain_document_chunks
+  end
+
+  private
+
+  def create_document_chunks_table
+    create_table :captain_document_chunks do |t|
+      t.bigint :document_id, null: false
+      t.bigint :assistant_id, null: false
+      t.bigint :account_id, null: false
+      t.text :content, null: false
+      t.text :context
+      t.vector :embedding, limit: 1536
+      t.integer :position, null: false
+      t.integer :token_count
+      t.column :searchable, :tsvector
+
+      t.timestamps
+    end
+
+    add_document_chunk_indexes
+  end
+
+  def add_document_chunk_indexes
+    add_index :captain_document_chunks, :document_id
+    add_index :captain_document_chunks, [:document_id, :position], unique: true
+    add_index :captain_document_chunks, [:account_id, :assistant_id]
+    add_index :captain_document_chunks, :searchable, using: :gin
+    add_index :captain_document_chunks,
+              :embedding,
+              using: :ivfflat,
+              name: 'index_captain_document_chunks_on_embedding',
+              opclass: :vector_l2_ops
+  end
+
+  def add_chunking_progress_columns
+    add_column :captain_documents, :chunking_status, :integer, default: 0, null: false
+    add_column :captain_documents, :expected_chunk_count, :integer, default: 0, null: false
+    add_column :captain_documents, :indexed_chunk_count, :integer, default: 0, null: false
+    add_column :captain_documents, :chunks_generated_at, :datetime
+    add_column :captain_documents, :last_chunk_error, :text
+
+    add_index :captain_documents, :chunking_status
+    add_index :captain_documents, :chunks_generated_at
+  end
+
+  def remove_chunking_progress_columns
+    remove_index :captain_documents, :chunking_status
+    remove_index :captain_documents, :chunks_generated_at
+
+    remove_column :captain_documents, :chunking_status
+    remove_column :captain_documents, :expected_chunk_count
+    remove_column :captain_documents, :indexed_chunk_count
+    remove_column :captain_documents, :chunks_generated_at
+    remove_column :captain_documents, :last_chunk_error
+  end
+end

enterprise/app/models/captain/assistant.rb

@@ -25,6 +25,7 @@ class Captain::Assistant < ApplicationRecord
 
   belongs_to :account
   has_many :documents, class_name: 'Captain::Document', dependent: :destroy_async
+  has_many :document_chunks, class_name: 'Captain::DocumentChunk', dependent: :destroy_async
   has_many :responses, class_name: 'Captain::AssistantResponse', dependent: :destroy_async
   has_many :captain_inboxes,
            class_name: 'CaptainInbox',

enterprise/app/models/captain/document.rb

@@ -3,8 +3,13 @@
 # Table name: captain_documents
 #
 #  id            :bigint           not null, primary key
+#  chunking_status      :integer          default("pending"), not null
+#  chunks_generated_at :datetime
 #  content       :text
+#  expected_chunk_count :integer          default(0), not null
 #  external_link :string           not null
+#  indexed_chunk_count :integer          default(0), not null
+#  last_chunk_error    :text
 #  metadata      :jsonb
 #  name          :string
 #  status        :integer          default("in_progress"), not null
@@ -18,6 +23,8 @@
 #  index_captain_documents_on_account_id                      (account_id)
 #  index_captain_documents_on_assistant_id                    (assistant_id)
 #  index_captain_documents_on_assistant_id_and_external_link  (assistant_id,external_link) UNIQUE
+#  index_captain_documents_on_chunks_generated_at             (chunks_generated_at)
+#  index_captain_documents_on_chunking_status                 (chunking_status)
 #  index_captain_documents_on_status                          (status)
 #
 class Captain::Document < ApplicationRecord
@@ -26,6 +33,7 @@ class Captain::Document < ApplicationRecord
 
   belongs_to :assistant, class_name: 'Captain::Assistant'
   has_many :responses, class_name: 'Captain::AssistantResponse', dependent: :destroy, as: :documentable
+  has_many :chunks, class_name: 'Captain::DocumentChunk', dependent: :destroy
   belongs_to :account
   has_one_attached :pdf_file
 
@@ -44,6 +52,14 @@ class Captain::Document < ApplicationRecord
     available: 1
   }
 
+  enum chunking_status: {
+    pending: 0,
+    chunking: 1,
+    indexing: 2,
+    ready: 3,
+    failed: 4
+  }, _prefix: true
+
   before_create :ensure_within_plan_limit
   after_create_commit :enqueue_crawl_job
   after_create_commit :update_document_usage

enterprise/app/models/captain/document_chunk.rb

@@ -0,0 +1,38 @@
+# == Schema Information
+#
+# Table name: captain_document_chunks
+#
+#  id           :bigint           not null, primary key
+#  content      :text             not null
+#  context      :text
+#  embedding    :vector(1536)
+#  position     :integer          not null
+#  searchable   :tsvector
+#  token_count  :integer
+#  created_at   :datetime         not null
+#  updated_at   :datetime         not null
+#  account_id   :bigint           not null
+#  assistant_id :bigint           not null
+#  document_id  :bigint           not null
+#
+# Indexes
+#
+#  index_captain_document_chunks_on_account_id_and_assistant_id  (account_id,assistant_id)
+#  index_captain_document_chunks_on_document_id                  (document_id)
+#  index_captain_document_chunks_on_document_id_and_position     (document_id,position) UNIQUE
+#  index_captain_document_chunks_on_embedding                    (embedding) USING ivfflat
+#  index_captain_document_chunks_on_searchable                   (searchable) USING gin
+#
+class Captain::DocumentChunk < ApplicationRecord
+  self.table_name = 'captain_document_chunks'
+
+  belongs_to :document, class_name: 'Captain::Document'
+  belongs_to :assistant, class_name: 'Captain::Assistant'
+  belongs_to :account
+
+  has_neighbors :embedding, normalize: true
+
+  validates :content, presence: true
+  validates :position, presence: true
+  validates :position, uniqueness: { scope: :document_id }
+end

enterprise/app/models/enterprise/concerns/account.rb

@@ -12,6 +12,7 @@ module Enterprise::Concerns::Account
     has_many :captain_assistants, dependent: :destroy_async, class_name: 'Captain::Assistant'
     has_many :captain_assistant_responses, dependent: :destroy_async, class_name: 'Captain::AssistantResponse'
     has_many :captain_documents, dependent: :destroy_async, class_name: 'Captain::Document'
+    has_many :captain_document_chunks, dependent: :destroy_async, class_name: 'Captain::DocumentChunk'
     has_many :captain_custom_tools, dependent: :destroy_async, class_name: 'Captain::CustomTool'
 
     has_many :copilot_threads, dependent: :destroy_async

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

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

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

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

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

@@ -0,0 +1,85 @@
+class Captain::Documents::ContextGenerationService < Llm::BaseAiService
+  include Integrations::LlmInstrumentation
+
+  MAX_DOCUMENT_CHARACTERS = 20_000
+  MAX_CHUNK_CHARACTERS = 6_000
+  DEFAULT_MODEL = 'gpt-4.1'.freeze
+
+  def initialize(document_content:, chunk_content:, account_id:, model: nil)
+    super()
+    @document_content = document_content.to_s
+    @chunk_content = chunk_content.to_s
+    @account_id = account_id
+    @model = resolve_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
+
+  def resolve_model(model)
+    return model if model.present?
+
+    context_model = InstallationConfig.find_by(name: 'CAPTAIN_CONTEXT_GENERATION_MODEL')&.value.presence
+    open_ai_model = InstallationConfig.find_by(name: 'CAPTAIN_OPEN_AI_MODEL')&.value.presence
+    context_model || open_ai_model || DEFAULT_MODEL
+  end
+end

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

@@ -0,0 +1,53 @@
+require 'rails_helper'
+
+RSpec.describe Captain::Documents::ChunkEmbeddingService do
+  let(:embedding_service) { instance_double(Captain::Llm::EmbeddingService) }
+
+  before do
+    allow(Captain::Llm::EmbeddingService).to receive(:new).with(account_id: 7).and_return(embedding_service)
+    allow(embedding_service).to receive(:get_embedding).and_return([0.1, 0.2, 0.3])
+  end
+
+  describe '#embed' do
+    it 'embeds combined context and content when both are present' do
+      service = described_class.new(account_id: 7)
+
+      result = service.embed(content: 'Starts at $19', context: 'Pricing page')
+
+      expect(result).to eq([0.1, 0.2, 0.3])
+      expect(embedding_service).to have_received(:get_embedding).with("Pricing page\n\nStarts at $19")
+    end
+
+    it 'returns empty embedding when both content and context are blank' do
+      service = described_class.new(account_id: 7)
+
+      result = service.embed(content: '', context: nil)
+
+      expect(result).to eq([])
+      expect(embedding_service).not_to have_received(:get_embedding)
+    end
+  end
+
+  describe '#build_record_attributes' do
+    it 'returns attributes ready for persisting a document chunk record' do
+      document = instance_double(Captain::Document, id: 10, assistant_id: 11, account_id: 7)
+      chunk = { position: 3, content: 'Chunk body', token_count: 42 }
+
+      service = described_class.new(account_id: 7)
+      result = service.build_record_attributes(document: document, chunk: chunk, context: 'Chunk context')
+
+      expect(result).to include(
+        document_id: 10,
+        assistant_id: 11,
+        account_id: 7,
+        position: 3,
+        content: 'Chunk body',
+        token_count: 42,
+        context: 'Chunk context',
+        embedding: [0.1, 0.2, 0.3]
+      )
+      expect(result[:created_at]).to be_present
+      expect(result[:updated_at]).to be_present
+    end
+  end
+end

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

@@ -0,0 +1,57 @@
+require 'rails_helper'
+
+RSpec.describe Captain::Documents::ChunkingService do
+  describe '#chunk' do
+    it 'returns no chunks for blank content' do
+      result = described_class.new('').chunk
+
+      expect(result).to eq([])
+    end
+
+    it 'returns ordered chunks with position and token count' do
+      content = <<~TEXT
+        # Pricing
+        Starts at $19 per agent.
+
+        ## Limits
+        Includes 100 contacts.
+
+        ## Support
+        Email support is available 24/7.
+      TEXT
+
+      result = described_class.new(
+        content,
+        target_tokens: 8,
+        min_tokens: 4,
+        max_tokens: 10,
+        overlap_tokens: 0
+      ).chunk
+
+      expect(result.size).to be >= 2
+      expect(result.map { |chunk| chunk[:position] }).to eq((0...result.size).to_a)
+      expect(result).to all(include(:content, :token_count, :position))
+      expect(result).to all(satisfy { |chunk| chunk[:token_count].positive? })
+      expect(result.first[:content]).to include('Section: Pricing')
+    end
+
+    it 'adds overlap from previous chunk to the next chunk' do
+      content = <<~TEXT
+        one two three four five six seven eight nine ten
+
+        alpha beta gamma delta epsilon zeta eta theta iota kappa
+      TEXT
+
+      result = described_class.new(
+        content,
+        target_tokens: 5,
+        min_tokens: 3,
+        max_tokens: 5,
+        overlap_tokens: 2
+      ).chunk
+
+      expect(result.size).to eq(2)
+      expect(result.last[:content]).to include('nine ten')
+    end
+  end
+end

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

@@ -0,0 +1,108 @@
+require 'rails_helper'
+require 'ostruct'
+
+RSpec.describe Captain::Documents::ContextGenerationService do
+  let(:chat) { instance_double(RubyLLM::Chat) }
+  let(:response) { instance_double(RubyLLM::Message, content: "  Pricing page context.  \n") }
+
+  before do
+    InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_OPEN_AI_API_KEY').update!(value: 'test-key')
+    InstallationConfig.where(name: %w[CAPTAIN_CONTEXT_GENERATION_MODEL CAPTAIN_OPEN_AI_MODEL]).delete_all
+
+    allow(RubyLLM).to receive(:chat).and_return(chat)
+    allow(chat).to receive(:with_temperature).and_return(chat)
+    allow(chat).to receive(:with_instructions).and_return(chat)
+    allow(chat).to receive(:ask).and_return(response)
+  end
+
+  describe '#generate' do
+    it 'returns stripped context text from the model response' do
+      service = described_class.new(
+        document_content: 'Pricing docs for all plans',
+        chunk_content: 'Business plan starts at $19/agent/month.',
+        account_id: 1
+      )
+
+      result = service.generate
+
+      expect(result).to eq('Pricing page context.')
+    end
+
+    it 'uses explicit model when provided' do
+      service = described_class.new(
+        document_content: 'Doc text',
+        chunk_content: 'Chunk text',
+        account_id: 1,
+        model: 'gpt-4.1-mini'
+      )
+
+      expect(RubyLLM).to receive(:chat).with(model: 'gpt-4.1-mini').and_return(chat)
+
+      service.generate
+    end
+
+    it 'uses CAPTAIN_CONTEXT_GENERATION_MODEL when configured' do
+      create(:installation_config, name: 'CAPTAIN_CONTEXT_GENERATION_MODEL', value: 'gpt-4.1-mini')
+
+      service = described_class.new(
+        document_content: 'Doc text',
+        chunk_content: 'Chunk text',
+        account_id: 1
+      )
+
+      expect(RubyLLM).to receive(:chat).with(model: 'gpt-4.1-mini').and_return(chat)
+
+      service.generate
+    end
+
+    it 'uses CAPTAIN_OPEN_AI_MODEL when context model is not configured' do
+      create(:installation_config, name: 'CAPTAIN_OPEN_AI_MODEL', value: 'gpt-4o-mini')
+
+      service = described_class.new(
+        document_content: 'Doc text',
+        chunk_content: 'Chunk text',
+        account_id: 1
+      )
+
+      expect(RubyLLM).to receive(:chat).with(model: 'gpt-4o-mini').and_return(chat)
+
+      service.generate
+    end
+
+    it 'falls back to the service default model when no model config exists' do
+      service = described_class.new(
+        document_content: 'Doc text',
+        chunk_content: 'Chunk text',
+        account_id: 1
+      )
+
+      expect(RubyLLM).to receive(:chat).with(model: 'gpt-4.1').and_return(chat)
+
+      service.generate
+    end
+
+    it 'returns empty string when chunk content is blank' do
+      service = described_class.new(
+        document_content: 'Doc text',
+        chunk_content: '',
+        account_id: 1
+      )
+
+      expect(service.generate).to eq('')
+      expect(RubyLLM).not_to have_received(:chat)
+    end
+
+    it 'returns empty string when LLM call fails' do
+      error = RubyLLM::Error.new(OpenStruct.new(body: {}), 'request failed')
+      allow(chat).to receive(:ask).and_raise(error)
+
+      service = described_class.new(
+        document_content: 'Doc text',
+        chunk_content: 'Chunk text',
+        account_id: 1
+      )
+
+      expect(service.generate).to eq('')
+    end
+  end
+end