Compare commits
7
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1f352b8143 | ||
|
|
0c03c2bb4d | ||
|
|
ed19b29a3c | ||
|
|
a96d763168 | ||
|
|
15eb6b0e3f | ||
|
|
7c7200d98e | ||
|
|
d3c91a29ca |
@@ -199,6 +199,16 @@
|
||||
display_title: 'Embedding Model (optional)'
|
||||
description: 'The embedding model configured for use in Captain AI. Default: text-embedding-3-small'
|
||||
locked: false
|
||||
- name: CAPTAIN_CONTEXT_GENERATION_MODEL
|
||||
display_title: 'Context Generation Model (optional)'
|
||||
description: 'Model used to generate contextual prefixes for knowledge chunks. Default: gpt-4.1'
|
||||
locked: false
|
||||
- name: CAPTAIN_DOCUMENT_CHUNKING_ENABLED
|
||||
display_title: 'Enable Document Chunking (v1)'
|
||||
description: 'Enables parallel chunk generation for Captain documents. Keep disabled until v1 rollout readiness.'
|
||||
value: false
|
||||
locked: false
|
||||
type: boolean
|
||||
- name: CAPTAIN_FIRECRAWL_API_KEY
|
||||
display_title: 'FireCrawl API Key (optional)'
|
||||
description: 'The FireCrawl API key for the Captain AI service'
|
||||
|
||||
@@ -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
|
||||
@@ -51,6 +51,7 @@ class Api::V1::Accounts::Captain::AssistantsController < Api::V1::Accounts::Base
|
||||
permitted = params.require(:assistant).permit(:name, :description,
|
||||
config: [
|
||||
:product_name, :feature_faq, :feature_memory, :feature_citation,
|
||||
:feature_document_faq_generation,
|
||||
:welcome_message, :handoff_message, :resolution_message,
|
||||
:instructions, :temperature
|
||||
])
|
||||
|
||||
@@ -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',
|
||||
@@ -36,7 +37,7 @@ class Captain::Assistant < ApplicationRecord
|
||||
has_many :copilot_threads, dependent: :destroy_async
|
||||
has_many :scenarios, class_name: 'Captain::Scenario', dependent: :destroy_async
|
||||
|
||||
store_accessor :config, :temperature, :feature_faq, :feature_memory, :product_name
|
||||
store_accessor :config, :temperature, :feature_faq, :feature_memory, :feature_document_faq_generation, :product_name
|
||||
|
||||
validates :name, presence: true
|
||||
validates :description, presence: true
|
||||
|
||||
@@ -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,11 +52,20 @@ 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
|
||||
after_destroy :update_document_usage
|
||||
after_commit :enqueue_response_builder_job
|
||||
after_commit :enqueue_chunk_builder_job
|
||||
scope :ordered, -> { order(created_at: :desc) }
|
||||
|
||||
scope :for_account, ->(account_id) { where(account_id: account_id) }
|
||||
@@ -100,15 +117,41 @@ class Captain::Document < ApplicationRecord
|
||||
Captain::Documents::ResponseBuilderJob.perform_later(self)
|
||||
end
|
||||
|
||||
def enqueue_chunk_builder_job
|
||||
return unless should_enqueue_chunk_builder_job?
|
||||
|
||||
Captain::Documents::ChunkBuilderJob.perform_later(self)
|
||||
end
|
||||
|
||||
def should_enqueue_response_builder?
|
||||
return false if destroyed?
|
||||
return false unless available?
|
||||
return false unless document_faq_generation_enabled?
|
||||
|
||||
return saved_change_to_status? if pdf_document?
|
||||
|
||||
(saved_change_to_status? || saved_change_to_content?) && content.present?
|
||||
end
|
||||
|
||||
def should_enqueue_chunk_builder_job?
|
||||
return false unless chunk_builder_enabled?
|
||||
return false if destroyed?
|
||||
return false unless available?
|
||||
return false if pdf_document?
|
||||
|
||||
(saved_change_to_status? || saved_change_to_content?) && content.present?
|
||||
end
|
||||
|
||||
def chunk_builder_enabled?
|
||||
value = InstallationConfig.find_by(name: 'CAPTAIN_DOCUMENT_CHUNKING_ENABLED')&.value
|
||||
ActiveModel::Type::Boolean.new.cast(value)
|
||||
end
|
||||
|
||||
def document_faq_generation_enabled?
|
||||
value = assistant&.config&.fetch('feature_document_faq_generation', true)
|
||||
ActiveModel::Type::Boolean.new.cast(value)
|
||||
end
|
||||
|
||||
def update_document_usage
|
||||
account.update_document_usage
|
||||
end
|
||||
|
||||
@@ -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
|
||||
@@ -69,7 +69,8 @@ RSpec.describe 'Api::V1::Accounts::Captain::Assistants', type: :request do
|
||||
product_name: 'Chatwoot',
|
||||
feature_faq: true,
|
||||
feature_memory: false,
|
||||
feature_citation: true
|
||||
feature_citation: true,
|
||||
feature_document_faq_generation: false
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -111,6 +112,18 @@ RSpec.describe 'Api::V1::Accounts::Captain::Assistants', type: :request do
|
||||
expect(response).to have_http_status(:success)
|
||||
end
|
||||
|
||||
it 'creates an assistant with feature_document_faq_generation disabled' do
|
||||
expect do
|
||||
post "/api/v1/accounts/#{account.id}/captain/assistants",
|
||||
params: valid_attributes,
|
||||
headers: admin.create_new_auth_token,
|
||||
as: :json
|
||||
end.to change(Captain::Assistant, :count).by(1)
|
||||
|
||||
expect(response).to have_http_status(:success)
|
||||
expect(json_response[:config][:feature_document_faq_generation]).to be(false)
|
||||
end
|
||||
|
||||
it 'creates an assistant with feature_citation disabled' do
|
||||
attributes_with_disabled_citation = valid_attributes.deep_dup
|
||||
attributes_with_disabled_citation[:assistant][:config][:feature_citation] = false
|
||||
@@ -216,6 +229,18 @@ RSpec.describe 'Api::V1::Accounts::Captain::Assistants', type: :request do
|
||||
expect(response).to have_http_status(:success)
|
||||
expect(json_response[:config][:feature_citation]).to be(false)
|
||||
end
|
||||
|
||||
it 'updates feature_document_faq_generation config' do
|
||||
assistant.update!(config: { 'feature_document_faq_generation' => true })
|
||||
|
||||
patch "/api/v1/accounts/#{account.id}/captain/assistants/#{assistant.id}",
|
||||
params: { assistant: { config: { feature_document_faq_generation: false } } },
|
||||
headers: admin.create_new_auth_token,
|
||||
as: :json
|
||||
|
||||
expect(response).to have_http_status(:success)
|
||||
expect(json_response[:config][:feature_document_faq_generation]).to be(false)
|
||||
end
|
||||
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
|
||||
@@ -193,6 +193,14 @@ RSpec.describe Captain::Document, type: :model do
|
||||
document.update!(metadata: { 'title' => 'Updated' })
|
||||
end.not_to have_enqueued_job(Captain::Documents::ResponseBuilderJob)
|
||||
end
|
||||
|
||||
it 'does not enqueue when document FAQ generation is disabled for the assistant' do
|
||||
assistant.update!(config: (assistant.config || {}).merge('feature_document_faq_generation' => false))
|
||||
|
||||
expect do
|
||||
create(:captain_document, assistant: assistant, account: account, status: :available, content: 'Doc content')
|
||||
end.not_to have_enqueued_job(Captain::Documents::ResponseBuilderJob)
|
||||
end
|
||||
end
|
||||
|
||||
describe 'PDF documents' do
|
||||
@@ -250,4 +258,53 @@ RSpec.describe Captain::Document, type: :model do
|
||||
end.not_to have_enqueued_job(Captain::Documents::ResponseBuilderJob)
|
||||
end
|
||||
end
|
||||
|
||||
describe 'chunk builder job callback' do
|
||||
before do
|
||||
clear_enqueued_jobs
|
||||
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_DOCUMENT_CHUNKING_ENABLED').update!(value: 'true')
|
||||
end
|
||||
|
||||
it 'enqueues for available non-PDF documents with content' do
|
||||
expect do
|
||||
create(:captain_document, assistant: assistant, account: account, status: :available, content: 'Doc content')
|
||||
end.to have_enqueued_job(Captain::Documents::ChunkBuilderJob)
|
||||
end
|
||||
|
||||
it 'does not enqueue when chunking flag is disabled' do
|
||||
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_DOCUMENT_CHUNKING_ENABLED').update!(value: 'false')
|
||||
|
||||
expect do
|
||||
create(:captain_document, assistant: assistant, account: account, status: :available, content: 'Doc content')
|
||||
end.not_to have_enqueued_job(Captain::Documents::ChunkBuilderJob)
|
||||
end
|
||||
|
||||
it 'does not enqueue for PDF documents' do
|
||||
document = build(:captain_document, assistant: assistant, account: account, status: :available, content: nil)
|
||||
document.pdf_file.attach(
|
||||
io: StringIO.new('PDF content'),
|
||||
filename: 'sample.pdf',
|
||||
content_type: 'application/pdf'
|
||||
)
|
||||
|
||||
expect do
|
||||
document.save!
|
||||
end.not_to have_enqueued_job(Captain::Documents::ChunkBuilderJob)
|
||||
end
|
||||
|
||||
it 'enqueues when content is updated on available non-PDF document' do
|
||||
document = create(
|
||||
:captain_document,
|
||||
assistant: assistant,
|
||||
account: account,
|
||||
status: :available,
|
||||
content: nil
|
||||
)
|
||||
clear_enqueued_jobs
|
||||
|
||||
expect do
|
||||
document.update!(content: 'Updated document content')
|
||||
end.to have_enqueued_job(Captain::Documents::ChunkBuilderJob)
|
||||
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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user