Compare commits

...
36 changed files with 1969 additions and 16 deletions
+10
View File
@@ -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
])
@@ -1,6 +1,7 @@
module Captain::ChatGenerationRecorder
extend ActiveSupport::Concern
include Integrations::LlmInstrumentationConstants
include Integrations::LlmUsageDetailsBuilder
private
@@ -30,12 +31,16 @@ module Captain::ChatGenerationRecorder
end
def generation_attributes(chat, message)
usage_details = generation_usage_details(message)
{
ATTR_GEN_AI_PROVIDER => determine_provider(model),
ATTR_GEN_AI_REQUEST_MODEL => model,
ATTR_GEN_AI_REQUEST_TEMPERATURE => temperature,
ATTR_GEN_AI_USAGE_INPUT_TOKENS => message.input_tokens,
ATTR_GEN_AI_USAGE_INPUT_TOKENS => usage_details[:input],
ATTR_GEN_AI_USAGE_OUTPUT_TOKENS => message.respond_to?(:output_tokens) ? message.output_tokens : nil,
ATTR_GEN_AI_USAGE_TOTAL_TOKENS => usage_details[:total],
ATTR_LANGFUSE_OBSERVATION_USAGE_DETAILS => usage_details.to_json.presence,
ATTR_LANGFUSE_OBSERVATION_INPUT => format_input_messages(chat),
ATTR_LANGFUSE_OBSERVATION_OUTPUT => message.respond_to?(:content) ? message.content.to_s : nil
}
@@ -44,4 +49,8 @@ module Captain::ChatGenerationRecorder
def format_input_messages(chat)
chat.messages[0...-1].map { |m| { role: m.role.to_s, content: m.content.to_s } }.to_json
end
def generation_usage_details(message)
usage_details_from_message(message, provider: determine_provider(model))
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
+2 -1
View File
@@ -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
+43
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,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,109 @@
class Captain::Documents::ChunkRerankerService
include Integrations::LlmInstrumentation
DEFAULT_MODEL = 'rerank-v4.0-pro'.freeze
DEFAULT_ENDPOINT = 'https://api.cohere.com/v2/rerank'.freeze
MAX_RERANK_TEXT_CHARACTERS = 1_200
def initialize(account_id:)
@account_id = account_id
@model = resolve_model
@api_key = InstallationConfig.find_by(name: 'CAPTAIN_COHERE_API_KEY')&.value.to_s
@endpoint =
InstallationConfig.find_by(name: 'CAPTAIN_COHERE_RERANK_ENDPOINT')&.value.to_s.presence ||
InstallationConfig.find_by(name: 'CAPTAIN_COHERE_API_BASE')&.value.to_s.presence ||
DEFAULT_ENDPOINT
end
def rerank(query:, candidates:, limit:)
return candidates.first(limit) if candidates.blank?
return candidates.first(limit) if @api_key.blank?
reranked_indices = fetch_reranked_indices(query, candidates)
return candidates.first(limit) if reranked_indices.blank?
reranked_chunks = reranked_indices.filter_map { |index| candidates[index] }
remaining_chunks = candidates.reject { |chunk| reranked_chunks.include?(chunk) }
(reranked_chunks + remaining_chunks).first(limit)
rescue StandardError => e
Rails.logger.warn "Chunk reranker failed: #{e.message}"
candidates.first(limit)
end
private
def reranker_text(chunk)
[chunk.context, chunk.content]
.compact
.join("\n\n")
.squish
.first(MAX_RERANK_TEXT_CHARACTERS)
end
def instrumentation_params(query, documents)
{
span_name: 'llm.captain.chunk_rerank',
model: @model,
feature_name: 'chunk_rerank',
account_id: @account_id,
messages: [
{ role: 'user', content: { query: query, document_count: documents.size }.to_json }
]
}
end
def resolve_model
InstallationConfig.find_by(name: 'CAPTAIN_CHUNK_RERANK_MODEL')&.value.presence || DEFAULT_MODEL
end
def fetch_reranked_indices(query, candidates)
documents = candidates.map { |chunk| reranker_text(chunk) }
response_payload = instrument_llm_call(instrumentation_params(query, documents)) do
cohere_rerank(query, documents)
end
parse_reranked_indices(response_payload, candidates.size)
end
def cohere_rerank(query, documents)
response = HTTParty.post(
@endpoint,
body: {
model: @model,
query: query,
documents: documents,
top_n: documents.size
}.to_json,
headers: cohere_headers
)
raise "Cohere rerank request failed with status #{response.code}" unless response.success?
parsed = response.parsed_response
raise 'Cohere rerank response is invalid' unless parsed.is_a?(Hash)
parsed
end
def parse_reranked_indices(payload, candidate_count)
results = payload['results']
return [] unless results.is_a?(Array)
results
.filter_map do |result|
index = Integer(result['index'], exception: false)
next nil if index.blank?
next nil if index.negative? || index >= candidate_count
index
end
.uniq
end
def cohere_headers
{
'Authorization' => "Bearer #{@api_key}",
'Content-Type' => 'application/json'
}
end
end
@@ -0,0 +1,217 @@
class Captain::Documents::ChunkingService
DEFAULT_TARGET_TOKENS = 600
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
BOILERPLATE_LINE_PATTERNS = [
/help center home page/i,
/theming_assets/i,
%r{hc/change_language/}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)
@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)
cleaned_content = remove_boilerplate_sections(remove_boilerplate_lines(content))
heading_path = []
cleaned_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 remove_boilerplate_sections(content)
content
.split(/\n{2,}/)
.map(&:strip)
.reject(&:blank?)
.reject { |section| boilerplate_section?(section) }
.join("\n\n")
end
def remove_boilerplate_lines(content)
content
.lines
.reject { |line| boilerplate_line?(line) }
.join
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 boilerplate_line?(line)
normalized = line.to_s.downcase.strip
return false if normalized.blank?
return true if BOILERPLATE_LINE_PATTERNS.any? { |pattern| normalized.match?(pattern) }
markdown_links = normalized.scan(/\[[^\]]+\]\([^)]+\)/).size
return true if normalized.include?('change_language/') && markdown_links >= 3
return false unless markdown_links >= 6
non_link_tokens = normalized
.gsub(/\[[^\]]+\]\([^)]+\)/, ' ')
.gsub(/[^a-z0-9\s]/, ' ')
.squeeze(' ')
.strip
.split
non_link_tokens.length <= 12
end
def link_heavy_navigation_section?(section)
lines = section.lines.map(&:strip).reject(&:blank?)
return false if lines.size < 3
markdown_links = section.scan(/\[[^\]]+\]\([^)]+\)/).size
linked_lines = lines.count { |line| link_line?(line) }
return false unless markdown_links >= 2 && linked_lines >= 3
section.scan(/\b[\w']+\b/).size <= 180
end
def link_line?(line) = line.match?(/\[[^\]]+\]\([^)]+\)/) || line.start_with?('* [', '- [')
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,87 @@
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 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
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,212 @@
class Captain::Documents::HybridChunkSearchService
DEFAULT_VECTOR_LIMIT = 20
DEFAULT_BM25_LIMIT = 20
DEFAULT_BM25_CANDIDATE_LIMIT = 200
DEFAULT_RRF_LIMIT = 20
DEFAULT_RESULT_LIMIT = 5
DEFAULT_RERANKER_ENABLED = true
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
@embedding_service = Captain::Llm::EmbeddingService.new(account_id: assistant.account_id)
@reranker = Captain::Documents::ChunkRerankerService.new(account_id: assistant.account_id)
end
def search(query, limit: DEFAULT_RESULT_LIMIT)
return [] if query.blank?
vector_results = vector_search(query)
bm25_results = bm25_search(query)
rrf_results = rank_results(vector_results, bm25_results, DEFAULT_RRF_LIMIT)
rerank_results(query, rrf_results, limit)
end
private
def base_scope
Captain::DocumentChunk
.where(account_id: @assistant.account_id, assistant_id: @assistant.id)
.joins(:document)
.merge(Captain::Document.chunking_status_ready)
.includes(:document)
end
def vector_search(query)
embedding = @embedding_service.get_embedding(query)
return [] if embedding.blank?
base_scope
.all
.nearest_neighbors(:embedding, embedding, distance: 'cosine')
.limit(DEFAULT_VECTOR_LIMIT)
.to_a
end
def bm25_search(query)
query_terms = bm25_terms(query).uniq
return [] if query_terms.blank?
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)
score_by_id = Hash.new(0.0)
vector_results.each_with_index do |chunk, index|
score_by_id[chunk.id] += 1.0 / (RRF_K + index + 1)
end
bm25_results.each_with_index do |chunk, index|
score_by_id[chunk.id] += 1.0 / (RRF_K + index + 1)
end
records_by_id = (vector_results + bm25_results).index_by(&:id)
score_by_id
.sort_by { |(_id, score)| -score }
.first(limit)
.filter_map { |(id, _score)| records_by_id[id] }
end
def rerank_results(query, candidates, limit)
return candidates.first(limit) unless reranking_enabled?
return candidates.first(limit) if candidates.blank?
reranked_results = @reranker.rerank(query: query, candidates: candidates, limit: limit)
return candidates.first(limit) if reranked_results.blank?
reranked_results
rescue StandardError => e
Rails.logger.warn "Hybrid chunk rerank failed: #{e.message}"
candidates.first(limit)
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
def reranking_enabled?
value = InstallationConfig.find_by(name: 'CAPTAIN_CHUNK_RERANKING_ENABLED')&.value
return DEFAULT_RERANKER_ENABLED if value.blank?
ActiveModel::Type::Boolean.new.cast(value)
end
end
@@ -5,20 +5,73 @@ class Captain::Tools::FaqLookupTool < Captain::Tools::BasePublicTool
def perform(_tool_context, query:)
log_tool_usage('searching', { query: query })
# Use existing vector search on approved responses
responses = @assistant.responses.approved.search(query).to_a
faq_results, chunk_results = search_knowledge(query)
total_results = faq_results.size + chunk_results.size
if responses.empty?
if total_results.zero?
log_tool_usage('no_results', { query: query })
"No relevant FAQs found for: #{query}"
else
log_tool_usage('found_results', { query: query, count: responses.size })
format_responses(responses)
log_tool_usage('found_results', { query: query, count: total_results })
"#{format_chunk_results(chunk_results)}#{format_responses(faq_results)}"
end
end
private
def search_knowledge(query)
if chunk_retrieval_mode?
[
search_non_document_faqs(query),
Captain::Documents::HybridChunkSearchService.new(assistant: @assistant).search(query)
]
else
[@assistant.responses.approved.search(query).to_a, []]
end
end
def search_non_document_faqs(query)
@assistant.responses
.approved
.where.not(documentable_type: 'Captain::Document')
.search(query)
.to_a
end
def chunk_retrieval_mode?
return false unless chunk_builder_enabled?
value = @assistant.config&.fetch('feature_document_faq_generation', true)
!ActiveModel::Type::Boolean.new.cast(value)
end
def chunk_builder_enabled?
value = InstallationConfig.find_by(name: 'CAPTAIN_DOCUMENT_CHUNKING_ENABLED')&.value
ActiveModel::Type::Boolean.new.cast(value)
end
def format_chunk_results(chunks)
chunks.map { |chunk| format_chunk(chunk) }.join
end
def format_chunk(chunk)
document = chunk.document
source_link = document.external_link if should_show_document_source?(document)
title = document.name.presence || document.external_link
formatted = "
Article: #{title}
Context: #{chunk.context}
Content: #{chunk.content}
"
if source_link.present?
formatted += "
Source: #{source_link}
"
end
formatted
end
def format_responses(responses)
responses.map { |response| format_response(response) }.join
end
@@ -45,4 +98,11 @@ class Captain::Tools::FaqLookupTool < Captain::Tools::BasePublicTool
external_link = response.documentable.external_link
!external_link.start_with?('PDF:')
end
def should_show_document_source?(document)
return false if document.blank?
return false if document.external_link.blank?
!document.external_link.start_with?('PDF:')
end
end
+2
View File
@@ -146,6 +146,8 @@ class Integrations::LlmBaseService
usage: {
'prompt_tokens' => response.input_tokens,
'completion_tokens' => response.output_tokens,
'cached_tokens' => response.respond_to?(:cached_tokens) ? response.cached_tokens : nil,
'cache_creation_tokens' => response.respond_to?(:cache_creation_tokens) ? response.cache_creation_tokens : nil,
'total_tokens' => (response.input_tokens || 0) + (response.output_tokens || 0)
},
request_messages: messages
+1 -1
View File
@@ -16,7 +16,7 @@ module Integrations::LlmInstrumentation
setup_span_attributes(span, params)
result = yield
executed = true
record_completion(span, result)
record_completion(span, result, params)
result
end
rescue StandardError => e
@@ -2,6 +2,7 @@
module Integrations::LlmInstrumentationCompletionHelpers
include Integrations::LlmInstrumentationConstants
include Integrations::LlmUsageDetailsBuilder
private
@@ -73,9 +74,13 @@ module Integrations::LlmInstrumentationCompletionHelpers
usage = result[:usage] || result['usage']
return if usage.blank?
span.set_attribute(ATTR_GEN_AI_USAGE_INPUT_TOKENS, usage['prompt_tokens']) if usage['prompt_tokens']
span.set_attribute(ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, usage['completion_tokens']) if usage['completion_tokens']
span.set_attribute(ATTR_GEN_AI_USAGE_TOTAL_TOKENS, usage['total_tokens']) if usage['total_tokens']
usage_details = usage_details_from_hash(usage)
return if usage_details.blank?
span.set_attribute(ATTR_GEN_AI_USAGE_INPUT_TOKENS, usage_details[:input]) if usage_details[:input]
span.set_attribute(ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, usage_details[:output]) if usage_details[:output]
span.set_attribute(ATTR_GEN_AI_USAGE_TOTAL_TOKENS, usage_details[:total]) if usage_details[:total]
span.set_attribute(ATTR_LANGFUSE_OBSERVATION_USAGE_DETAILS, usage_details.to_json)
end
def set_error_attributes(span, result)
@@ -85,4 +90,14 @@ module Integrations::LlmInstrumentationCompletionHelpers
span.set_attribute(ATTR_GEN_AI_RESPONSE_ERROR, error.to_json)
span.status = OpenTelemetry::Trace::Status.error(error.to_s.truncate(1000))
end
def set_message_usage_metrics(span, message, provider: nil)
usage_details = usage_details_from_message(message, provider: provider)
return if usage_details.blank?
span.set_attribute(ATTR_GEN_AI_USAGE_INPUT_TOKENS, usage_details[:input]) if usage_details[:input]
span.set_attribute(ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, usage_details[:output]) if usage_details[:output]
span.set_attribute(ATTR_GEN_AI_USAGE_TOTAL_TOKENS, usage_details[:total]) if usage_details[:total]
span.set_attribute(ATTR_LANGFUSE_OBSERVATION_USAGE_DETAILS, usage_details.to_json)
end
end
@@ -29,4 +29,5 @@ module Integrations::LlmInstrumentationConstants
ATTR_LANGFUSE_OBSERVATION_TYPE = 'langfuse.observation.type'
ATTR_LANGFUSE_OBSERVATION_INPUT = 'langfuse.observation.input'
ATTR_LANGFUSE_OBSERVATION_OUTPUT = 'langfuse.observation.output'
ATTR_LANGFUSE_OBSERVATION_USAGE_DETAILS = 'langfuse.observation.usage_details'
end
@@ -24,10 +24,11 @@ module Integrations::LlmInstrumentationHelpers
set_metadata_attributes(span, params)
end
def record_completion(span, result)
def record_completion(span, result, params)
if result.respond_to?(:content)
span.set_attribute(ATTR_GEN_AI_COMPLETION_ROLE, result.role.to_s) if result.respond_to?(:role)
span.set_attribute(ATTR_GEN_AI_COMPLETION_CONTENT, result.content.to_s)
set_message_usage_metrics(span, result, provider: determine_provider(params[:model]))
elsif result.is_a?(Hash)
set_completion_attributes(span, result)
end
+25 -2
View File
@@ -4,6 +4,7 @@ require 'opentelemetry_config'
module Integrations::LlmInstrumentationSpans
include Integrations::LlmInstrumentationConstants
include Integrations::LlmUsageDetailsBuilder
def tracer
@tracer ||= OpentelemetryConfig.tracer
@@ -86,7 +87,29 @@ module Integrations::LlmInstrumentationSpans
end
def set_llm_turn_usage_attributes(span, message)
span.set_attribute(ATTR_GEN_AI_USAGE_INPUT_TOKENS, message.input_tokens) if message.respond_to?(:input_tokens) && message.input_tokens
span.set_attribute(ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, message.output_tokens) if message.respond_to?(:output_tokens) && message.output_tokens
usage_details = llm_turn_usage_details(message)
return if usage_details.blank?
span.set_attribute(ATTR_GEN_AI_USAGE_INPUT_TOKENS, usage_details[:input]) if usage_details[:input]
span.set_attribute(ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, usage_details[:output]) if usage_details[:output]
span.set_attribute(ATTR_GEN_AI_USAGE_TOTAL_TOKENS, usage_details[:total]) if usage_details[:total]
span.set_attribute(ATTR_LANGFUSE_OBSERVATION_USAGE_DETAILS, usage_details.to_json)
end
def llm_turn_usage_details(message)
usage_details_from_message(message, provider: llm_turn_provider(message))
end
def llm_turn_provider(message)
model_name = message.respond_to?(:model_id) ? message.model_id : nil
return 'openai' if model_name.blank?
model = model_name.to_s.downcase
LlmConstants::PROVIDER_PREFIXES.each do |provider, prefixes|
return provider if prefixes.any? { |prefix| model.start_with?(prefix) }
end
'openai'
end
end
@@ -0,0 +1,110 @@
# frozen_string_literal: true
module Integrations::LlmUsageDetailsBuilder
private
def usage_details_from_hash(usage)
usage_hash = normalize_usage_hash(usage)
return {} if usage_hash.blank?
usage_details_from_hash_values(usage_hash)
end
def usage_details_from_hash_values(usage_hash)
input_tokens_source = usage_hash.key?('input_tokens') ? 'input_tokens' : 'prompt_tokens'
cache_read_tokens = usage_cache_read_tokens(usage_hash)
cache_creation_tokens = usage_cache_creation_tokens(usage_hash)
input_tokens = normalized_uncached_input_tokens(
usage_hash[input_tokens_source],
cache_read_tokens,
input_tokens_source: input_tokens_source
)
output_tokens = usage_hash['output_tokens'] || usage_hash['completion_tokens']
total_tokens = usage_total_tokens(
usage_hash['total_tokens'],
input_tokens,
output_tokens,
cache_read_tokens,
cache_creation_tokens
)
build_usage_details(input_tokens, output_tokens, total_tokens, cache_read_tokens, cache_creation_tokens)
end
def usage_details_from_message(message, provider:)
input_tokens = message_token(message, :input_tokens)
output_tokens = message_token(message, :output_tokens)
cache_read_tokens = message_token(message, :cached_tokens)
cache_creation_tokens = message_token(message, :cache_creation_tokens)
input_tokens = normalized_uncached_input_tokens(
input_tokens,
cache_read_tokens,
provider: provider,
input_tokens_source: 'input_tokens'
)
total_tokens = total_tokens_from_parts(input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens)
build_usage_details(input_tokens, output_tokens, total_tokens, cache_read_tokens, cache_creation_tokens)
end
def build_usage_details(input_tokens, output_tokens, total_tokens, cache_read_tokens, cache_creation_tokens)
compact_usage_details(
input: input_tokens,
output: output_tokens,
total: total_tokens,
cache_read_input_tokens: cache_read_tokens,
cache_creation_input_tokens: cache_creation_tokens
)
end
def normalize_usage_hash(usage)
return usage.deep_stringify_keys if usage.respond_to?(:deep_stringify_keys)
return usage.to_h.transform_keys(&:to_s) if usage.respond_to?(:to_h)
{}
end
def usage_cache_read_tokens(usage_hash)
usage_hash['cache_read_input_tokens'] ||
usage_hash['cached_tokens'] ||
usage_hash.dig('prompt_tokens_details', 'cached_tokens')
end
def usage_cache_creation_tokens(usage_hash)
usage_hash['cache_creation_input_tokens'] ||
usage_hash['cache_creation_tokens'] ||
usage_hash['cache_creation']&.values&.compact&.sum
end
def usage_total_tokens(reported_total, input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens)
if cache_read_tokens || cache_creation_tokens
total_tokens_from_parts(input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens)
else
reported_total || total_tokens_from_parts(input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens)
end
end
def message_token(message, token_name)
message.respond_to?(token_name) ? message.public_send(token_name) : nil
end
def normalized_uncached_input_tokens(input_tokens, cache_read_tokens, provider: nil, input_tokens_source: nil)
return input_tokens if input_tokens.nil? || cache_read_tokens.nil?
if input_tokens_source == 'prompt_tokens' || provider == 'openai'
[input_tokens - cache_read_tokens, 0].max
else
input_tokens
end
end
def total_tokens_from_parts(input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens)
values = [input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens].compact
values.sum if values.present?
end
def compact_usage_details(details)
details.compact
end
end
+1
View File
@@ -9,6 +9,7 @@ module LlmConstants
PROVIDER_PREFIXES = {
'openai' => %w[gpt- o1 o3 o4 text-embedding- whisper- tts-],
'cohere' => %w[rerank- command- embed-],
'anthropic' => %w[claude-],
'google' => %w[gemini-],
'mistral' => %w[mistral- codestral-],
@@ -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
@@ -110,6 +110,48 @@ RSpec.describe Captain::Tools::FaqLookupTool, type: :model do
expect(result).to eq('No relevant FAQs found for: ')
end
end
context 'when chunk retrieval mode is enabled for assistant' do
let(:document) do
create(
:captain_document,
assistant: assistant,
account: account,
chunking_status: :ready,
status: :available,
external_link: 'https://help.example.com/pricing',
name: 'Pricing'
)
end
let(:chunk) do
create(
:captain_document_chunk,
document: document,
assistant: assistant,
account: account,
content: 'Business plan starts at $19.',
context: 'Pricing page details'
)
end
let(:chunk_search_service) { instance_double(Captain::Documents::HybridChunkSearchService) }
before do
assistant.update!(config: (assistant.config || {}).merge('feature_document_faq_generation' => false))
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_DOCUMENT_CHUNKING_ENABLED').update!(value: 'true')
allow(Captain::Documents::HybridChunkSearchService).to receive(:new).with(assistant: assistant).and_return(chunk_search_service)
allow(chunk_search_service).to receive(:search).with('pricing').and_return([chunk])
allow(Captain::AssistantResponse).to receive(:nearest_neighbors).and_return(Captain::AssistantResponse.none)
end
it 'returns chunk content for knowledge lookup' do
result = tool.perform(tool_context, query: 'pricing')
expect(result).to include('Article: Pricing')
expect(result).to include('Context: Pricing page details')
expect(result).to include('Content: Business plan starts at $19.')
expect(result).to include('Source: https://help.example.com/pricing')
end
end
end
describe '#active?' do
@@ -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,114 @@
require 'rails_helper'
RSpec.describe Captain::Documents::ChunkRerankerService do
let(:assistant) { create(:captain_assistant) }
let(:account) { assistant.account }
let(:service) { described_class.new(account_id: account.id) }
before do
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_COHERE_API_KEY').update!(value: 'cohere-test-key')
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_CHUNK_RERANK_MODEL').update!(value: 'rerank-v4.0-pro')
end
describe '#rerank' do
it 'returns chunks in reranked order when the model returns ids' do
document = create(:captain_document, account: account, assistant: assistant, chunking_status: :ready, status: :available)
first_chunk = create(
:captain_document_chunk,
document: document,
account: account,
assistant: assistant,
position: 0,
content: 'How to cancel a subscription'
)
second_chunk = create(
:captain_document_chunk,
document: document,
account: account,
assistant: assistant,
position: 1,
content: 'How to delete an account permanently'
)
response_payload = {
'results' => [
{ 'index' => 1, 'relevance_score' => 0.93 },
{ 'index' => 0, 'relevance_score' => 0.67 }
]
}
cohere_response = instance_double(
HTTParty::Response,
success?: true,
code: 200,
parsed_response: response_payload
)
allow(HTTParty).to receive(:post).and_return(cohere_response)
results = service.rerank(query: 'delete account', candidates: [first_chunk, second_chunk], limit: 2)
expect(results.map(&:id)).to eq([second_chunk.id, first_chunk.id])
end
it 'falls back to original order when response payload is invalid' do
document = create(:captain_document, account: account, assistant: assistant, chunking_status: :ready, status: :available)
first_chunk = create(
:captain_document_chunk,
document: document,
account: account,
assistant: assistant,
position: 0,
content: 'How to cancel a subscription'
)
second_chunk = create(
:captain_document_chunk,
document: document,
account: account,
assistant: assistant,
position: 1,
content: 'How to delete an account permanently'
)
cohere_response = instance_double(
HTTParty::Response,
success?: true,
code: 200,
parsed_response: { 'results' => nil }
)
allow(HTTParty).to receive(:post).and_return(cohere_response)
results = service.rerank(query: 'delete account', candidates: [first_chunk, second_chunk], limit: 2)
expect(results.map(&:id)).to eq([first_chunk.id, second_chunk.id])
end
it 'falls back to original order when API key is missing' do
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_COHERE_API_KEY').update!(value: '')
document = create(:captain_document, account: account, assistant: assistant, chunking_status: :ready, status: :available)
first_chunk = create(
:captain_document_chunk,
document: document,
account: account,
assistant: assistant,
position: 0,
content: 'How to cancel a subscription'
)
second_chunk = create(
:captain_document_chunk,
document: document,
account: account,
assistant: assistant,
position: 1,
content: 'How to delete an account permanently'
)
results = described_class.new(account_id: account.id).rerank(
query: 'delete account',
candidates: [first_chunk, second_chunk],
limit: 2
)
expect(results.map(&:id)).to eq([first_chunk.id, second_chunk.id])
end
end
end
@@ -0,0 +1,105 @@
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
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
it 'removes single-line navigation and language selector blobs before chunking' do
content = <<~TEXT
[![Bumble Support Help Center home page](https://support.bumble.com/hc/theming_assets/logo.svg)] [Date](https://bumble.com/en-us/date) [Friends](https://bumble.com/en-us/bff) [Bizz](https://bumble.com/en-us/bizz) [Safety](https://bumble.com/en-us/the-buzz/category/safety) [Dansk](https://support.bumble.com/hc/change_language/da?return_to=%2Fhc%2Fen-us) [Deutsch](https://support.bumble.com/hc/change_language/de?return_to=%2Fhc%2Fen-us) [Espanol](https://support.bumble.com/hc/change_language/es?return_to=%2Fhc%2Fen-us)
# How to delete my account
Open Settings and select Delete account.
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('Open Settings and select Delete account')
expect(combined_content).not_to include('Help Center home page')
expect(combined_content).not_to include('change_language')
expect(combined_content).not_to include('[Date](')
end
end
end
@@ -0,0 +1,121 @@
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 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',
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,182 @@
require 'rails_helper'
RSpec.describe Captain::Documents::HybridChunkSearchService do
let(:assistant) { create(:captain_assistant) }
let(:account) { assistant.account }
let(:service) { described_class.new(assistant: assistant) }
before do
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_OPEN_AI_API_KEY').update!(value: 'test-key')
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_CHUNK_RERANKING_ENABLED').update!(value: 'false')
end
describe '#search' do
it 'returns no results for blank query' do
expect(service.search('')).to eq([])
end
it 'returns chunk results from ready documents scoped to the assistant' do
ready_document = create(
:captain_document,
account: account,
assistant: assistant,
status: :available,
chunking_status: :ready
)
skipped_document = create(
:captain_document,
account: account,
assistant: assistant,
status: :available,
chunking_status: :pending
)
create(
:captain_document_chunk,
document: ready_document,
account: account,
assistant: assistant,
content: 'How to reset password for admins'
)
create(
:captain_document_chunk,
document: skipped_document,
account: account,
assistant: assistant,
content: 'How to reset password for members'
)
embedding_service = instance_double(Captain::Llm::EmbeddingService, get_embedding: [])
allow(Captain::Llm::EmbeddingService).to receive(:new).and_return(embedding_service)
results = service.search('reset password')
expect(results.size).to eq(1)
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
it 'reranks RRF candidates with the configured reranker' do
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_CHUNK_RERANKING_ENABLED').update!(value: 'true')
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_CHUNK_RERANK_MODEL').update!(value: 'rerank-v4.0-pro')
ready_document = create(
:captain_document,
account: account,
assistant: assistant,
status: :available,
chunking_status: :ready
)
first_chunk = create(
:captain_document_chunk,
document: ready_document,
account: account,
assistant: assistant,
position: 0,
content: 'Cancel subscription from billing settings on iOS and Android.'
)
second_chunk = create(
:captain_document_chunk,
document: ready_document,
account: account,
assistant: assistant,
position: 1,
content: 'Delete your account permanently from profile settings and confirm with password.'
)
embedding_service = instance_double(Captain::Llm::EmbeddingService, get_embedding: [])
allow(Captain::Llm::EmbeddingService).to receive(:new).and_return(embedding_service)
reranker_service = instance_double(
Captain::Documents::ChunkRerankerService,
rerank: [second_chunk, first_chunk]
)
allow(Captain::Documents::ChunkRerankerService).to receive(:new).and_return(reranker_service)
results = described_class.new(assistant: assistant).search('How do I cancel or delete my account?', limit: 2)
expect(results.map(&:id)).to eq([second_chunk.id, first_chunk.id])
end
it 'falls back to RRF order when reranking fails' do
ready_document = create(
:captain_document,
account: account,
assistant: assistant,
status: :available,
chunking_status: :ready
)
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.'
)
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)
fallback_results = described_class.new(assistant: assistant).search('how does incognito mode work', limit: 2)
InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_CHUNK_RERANKING_ENABLED').update!(value: 'true')
reranker_service = instance_double(Captain::Documents::ChunkRerankerService)
allow(reranker_service).to receive(:rerank).and_raise(StandardError, 'reranker unavailable')
allow(Captain::Documents::ChunkRerankerService).to receive(:new).and_return(reranker_service)
reranker_failed_results = described_class.new(assistant: assistant).search('how does incognito mode work', limit: 2)
expect(reranker_failed_results.map(&:id)).to eq(fallback_results.map(&:id))
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
@@ -176,6 +176,7 @@ RSpec.describe Integrations::LlmInstrumentation do
{
usage: {
'prompt_tokens' => 150,
'prompt_tokens_details' => { 'cached_tokens' => 20 },
'completion_tokens' => 200,
'total_tokens' => 350
}
@@ -183,9 +184,41 @@ RSpec.describe Integrations::LlmInstrumentation do
end
expect(result[:usage]['prompt_tokens']).to eq(150)
expect(mock_span).to have_received(:set_attribute).with('gen_ai.usage.input_tokens', 150)
expect(mock_span).to have_received(:set_attribute).with('gen_ai.usage.input_tokens', 130)
expect(mock_span).to have_received(:set_attribute).with('gen_ai.usage.output_tokens', 200)
expect(mock_span).to have_received(:set_attribute).with('gen_ai.usage.total_tokens', 350)
expect(mock_span).to have_received(:set_attribute)
.with('langfuse.observation.usage_details', '{"input":130,"output":200,"total":350,"cache_read_input_tokens":20}')
end
it 'sets usage metrics for RubyLLM message responses with cached tokens' do
mock_span = instance_double(OpenTelemetry::Trace::Span)
allow(mock_span).to receive(:set_attribute)
allow(mock_span).to receive(:status=)
mock_tracer = instance_double(OpenTelemetry::Trace::Tracer)
allow(instance).to receive(:tracer).and_return(mock_tracer)
allow(mock_tracer).to receive(:in_span).and_yield(mock_span)
llm_message = instance_double(
RubyLLM::Message,
role: :assistant,
content: 'AI response',
input_tokens: 150,
output_tokens: 50,
cached_tokens: 40,
cache_creation_tokens: nil
)
result = instance.instrument_llm_call(params) do
llm_message
end
expect(result).to eq(llm_message)
expect(mock_span).to have_received(:set_attribute).with('gen_ai.usage.input_tokens', 110)
expect(mock_span).to have_received(:set_attribute).with('gen_ai.usage.output_tokens', 50)
expect(mock_span).to have_received(:set_attribute).with('gen_ai.usage.total_tokens', 200)
expect(mock_span).to have_received(:set_attribute)
.with('langfuse.observation.usage_details', '{"input":110,"output":50,"total":200,"cache_read_input_tokens":40}')
end
it 'sets error attributes when result contains error' do