Compare commits

...
16 changed files with 783 additions and 0 deletions
@@ -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
@@ -0,0 +1,7 @@
class Captain::Documents::ChunkBuilderJob < ApplicationJob
queue_as :low
def perform(document)
Captain::Documents::ChunkBuilderService.new(document).process
end
end
@@ -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',
+16
View File
@@ -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
@@ -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
@@ -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
@@ -0,0 +1,56 @@
class Captain::Documents::ChunkBuilderService
def initialize(document)
@document = document
@embedding_service = Captain::Documents::ChunkEmbeddingService.new(account_id: document.account_id)
end
def process
raise ArgumentError, 'Document content is required for chunk building' if @document.content.blank?
chunks = Captain::Documents::ChunkingService.new(@document.content).chunk
@document.update!(
chunking_status: :chunking,
last_chunk_error: nil,
chunks_generated_at: nil,
expected_chunk_count: chunks.count,
indexed_chunk_count: 0
)
indexed_count = rebuild_chunks(chunks)
@document.update!(
chunking_status: :ready,
indexed_chunk_count: indexed_count,
chunks_generated_at: Time.current
)
rescue StandardError => e
@document.update!(chunking_status: :failed, last_chunk_error: e.message)
raise
end
private
def rebuild_chunks(chunks)
@document.update!(chunking_status: :indexing)
@document.chunks.delete_all
chunks.each_with_index do |chunk, index|
context = generate_context(chunk[:content])
attributes = @embedding_service.build_record_attributes(document: @document, chunk: chunk, context: context)
Captain::DocumentChunk.create!(attributes)
@document.update_column(:indexed_chunk_count, index + 1) # rubocop:disable Rails/SkipsModelValidations
end
chunks.count
end
def generate_context(chunk_content)
Captain::Documents::ContextGenerationService.new(
document_content: @document.content,
chunk_content: chunk_content,
account_id: @document.account_id
).generate
end
end
@@ -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,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
@@ -0,0 +1,25 @@
require 'rails_helper'
RSpec.describe Captain::Documents::ChunkBuilderJob, type: :job do
let(:document) { create(:captain_document, status: :available) }
let(:builder_service) { instance_double(Captain::Documents::ChunkBuilderService) }
describe '#perform' do
it 'runs the chunk builder service for the document' do
allow(Captain::Documents::ChunkBuilderService).to receive(:new).with(document).and_return(builder_service)
allow(builder_service).to receive(:process)
described_class.new.perform(document)
expect(Captain::Documents::ChunkBuilderService).to have_received(:new).with(document)
expect(builder_service).to have_received(:process)
end
it 'bubbles service errors for retry handling' do
allow(Captain::Documents::ChunkBuilderService).to receive(:new).with(document).and_return(builder_service)
allow(builder_service).to receive(:process).and_raise(StandardError, 'transient failure')
expect { described_class.new.perform(document) }.to raise_error(StandardError, 'transient failure')
end
end
end
@@ -0,0 +1,80 @@
require 'rails_helper'
RSpec.describe Captain::Documents::ChunkBuilderService do
let(:document) { create(:captain_document, status: :available, content: 'Pricing and limits content') }
let(:chunking_service) { instance_double(Captain::Documents::ChunkingService) }
let(:embedding_service) { instance_double(Captain::Documents::ChunkEmbeddingService) }
let(:context_service) { instance_double(Captain::Documents::ContextGenerationService) }
let(:chunks) do
[
{ position: 0, content: 'Chunk 1', token_count: 10 },
{ position: 1, content: 'Chunk 2', token_count: 12 }
]
end
let(:embedding_vector) { Array.new(1536, 0.1) }
before do
allow(Captain::Documents::ChunkingService).to receive(:new).and_return(chunking_service)
allow(chunking_service).to receive(:chunk).and_return(chunks)
allow(Captain::Documents::ChunkEmbeddingService).to receive(:new).and_return(embedding_service)
allow(embedding_service).to receive(:build_record_attributes) do |args|
chunk = args.fetch(:chunk)
{
document_id: document.id,
assistant_id: document.assistant_id,
account_id: document.account_id,
position: chunk.fetch(:position),
content: chunk.fetch(:content),
token_count: chunk[:token_count],
context: args[:context],
embedding: embedding_vector,
created_at: Time.current,
updated_at: Time.current
}
end
allow(Captain::Documents::ContextGenerationService).to receive(:new).and_return(context_service)
allow(context_service).to receive(:generate).and_return('Chunk context')
end
describe '#process' do
it 'builds chunks and marks the document ready with progress counters' do
described_class.new(document).process
document.reload
expect(document.chunking_status).to eq('ready')
expect(document.expected_chunk_count).to eq(2)
expect(document.indexed_chunk_count).to eq(2)
expect(document.chunks_generated_at).to be_present
expect(document.last_chunk_error).to be_nil
expect(document.chunks.count).to eq(2)
end
it 'replaces existing chunks when reprocessed (idempotent rebuild)' do
create(
:captain_document_chunk,
document: document,
assistant: document.assistant,
account: document.account,
position: 99,
embedding: nil
)
described_class.new(document).process
positions = document.chunks.reload.order(:position).pluck(:position)
expect(positions).to eq([0, 1])
end
it 'marks the document failed and stores the error when processing fails' do
allow(chunking_service).to receive(:chunk).and_raise(StandardError, 'chunking failed')
expect { described_class.new(document).process }.to raise_error(StandardError, 'chunking failed')
document.reload
expect(document.chunking_status).to eq('failed')
expect(document.last_chunk_error).to eq('chunking failed')
end
end
end
@@ -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
@@ -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
@@ -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
+12
View File
@@ -0,0 +1,12 @@
FactoryBot.define do
factory :captain_document_chunk, class: 'Captain::DocumentChunk' do
association :document, factory: :captain_document
assistant { document.assistant }
account { document.account }
content { 'Chunk content' }
add_attribute(:context) { 'Chunk context' }
embedding { Array.new(1536, 0.1) }
position { 0 }
token_count { 12 }
end
end