From 15eb6b0e3f4b0134b218242d95a7af88e90c2091 Mon Sep 17 00:00:00 2001 From: aakashb95 Date: Thu, 19 Feb 2026 14:49:32 +0530 Subject: [PATCH] test(captain): add behavior specs for chunk services --- .../documents/context_generation_service.rb | 14 ++- .../documents/chunk_embedding_service_spec.rb | 53 +++++++++ .../documents/chunking_service_spec.rb | 57 +++++++++ .../context_generation_service_spec.rb | 108 ++++++++++++++++++ 4 files changed, 229 insertions(+), 3 deletions(-) create mode 100644 spec/enterprise/services/captain/documents/chunk_embedding_service_spec.rb create mode 100644 spec/enterprise/services/captain/documents/chunking_service_spec.rb create mode 100644 spec/enterprise/services/captain/documents/context_generation_service_spec.rb diff --git a/enterprise/app/services/captain/documents/context_generation_service.rb b/enterprise/app/services/captain/documents/context_generation_service.rb index e4b53168b..fee5f1580 100644 --- a/enterprise/app/services/captain/documents/context_generation_service.rb +++ b/enterprise/app/services/captain/documents/context_generation_service.rb @@ -3,14 +3,14 @@ class Captain::Documents::ContextGenerationService < Llm::BaseAiService MAX_DOCUMENT_CHARACTERS = 20_000 MAX_CHUNK_CHARACTERS = 6_000 - DEFAULT_MODEL = 'gpt-4.1-mini'.freeze + DEFAULT_MODEL = 'gpt-4.1'.freeze - def initialize(document_content:, chunk_content:, account_id:, model: DEFAULT_MODEL) + 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 = model + @model = resolve_model(model) end def generate @@ -74,4 +74,12 @@ class Captain::Documents::ContextGenerationService < Llm::BaseAiService ] } 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 diff --git a/spec/enterprise/services/captain/documents/chunk_embedding_service_spec.rb b/spec/enterprise/services/captain/documents/chunk_embedding_service_spec.rb new file mode 100644 index 000000000..07f544e3a --- /dev/null +++ b/spec/enterprise/services/captain/documents/chunk_embedding_service_spec.rb @@ -0,0 +1,53 @@ +require 'rails_helper' + +RSpec.describe Captain::Documents::ChunkEmbeddingService do + let(:embedding_service) { instance_double(Captain::Llm::EmbeddingService) } + + before do + allow(Captain::Llm::EmbeddingService).to receive(:new).with(account_id: 7).and_return(embedding_service) + allow(embedding_service).to receive(:get_embedding).and_return([0.1, 0.2, 0.3]) + end + + describe '#embed' do + it 'embeds combined context and content when both are present' do + service = described_class.new(account_id: 7) + + result = service.embed(content: 'Starts at $19', context: 'Pricing page') + + expect(result).to eq([0.1, 0.2, 0.3]) + expect(embedding_service).to have_received(:get_embedding).with("Pricing page\n\nStarts at $19") + end + + it 'returns empty embedding when both content and context are blank' do + service = described_class.new(account_id: 7) + + result = service.embed(content: '', context: nil) + + expect(result).to eq([]) + expect(embedding_service).not_to have_received(:get_embedding) + end + end + + describe '#build_record_attributes' do + it 'returns attributes ready for persisting a document chunk record' do + document = instance_double(Captain::Document, id: 10, assistant_id: 11, account_id: 7) + chunk = { position: 3, content: 'Chunk body', token_count: 42 } + + service = described_class.new(account_id: 7) + result = service.build_record_attributes(document: document, chunk: chunk, context: 'Chunk context') + + expect(result).to include( + document_id: 10, + assistant_id: 11, + account_id: 7, + position: 3, + content: 'Chunk body', + token_count: 42, + context: 'Chunk context', + embedding: [0.1, 0.2, 0.3] + ) + expect(result[:created_at]).to be_present + expect(result[:updated_at]).to be_present + end + end +end diff --git a/spec/enterprise/services/captain/documents/chunking_service_spec.rb b/spec/enterprise/services/captain/documents/chunking_service_spec.rb new file mode 100644 index 000000000..72f2f61e3 --- /dev/null +++ b/spec/enterprise/services/captain/documents/chunking_service_spec.rb @@ -0,0 +1,57 @@ +require 'rails_helper' + +RSpec.describe Captain::Documents::ChunkingService do + describe '#chunk' do + it 'returns no chunks for blank content' do + result = described_class.new('').chunk + + expect(result).to eq([]) + end + + it 'returns ordered chunks with position and token count' do + content = <<~TEXT + # Pricing + Starts at $19 per agent. + + ## Limits + Includes 100 contacts. + + ## Support + Email support is available 24/7. + TEXT + + result = described_class.new( + content, + target_tokens: 8, + min_tokens: 4, + max_tokens: 10, + overlap_tokens: 0 + ).chunk + + expect(result.size).to be >= 2 + expect(result.map { |chunk| chunk[:position] }).to eq((0...result.size).to_a) + expect(result).to all(include(:content, :token_count, :position)) + expect(result).to all(satisfy { |chunk| chunk[:token_count].positive? }) + expect(result.first[:content]).to include('Section: Pricing') + end + + it 'adds overlap from previous chunk to the next chunk' do + content = <<~TEXT + one two three four five six seven eight nine ten + + alpha beta gamma delta epsilon zeta eta theta iota kappa + TEXT + + result = described_class.new( + content, + target_tokens: 5, + min_tokens: 3, + max_tokens: 5, + overlap_tokens: 2 + ).chunk + + expect(result.size).to eq(2) + expect(result.last[:content]).to include('nine ten') + end + end +end diff --git a/spec/enterprise/services/captain/documents/context_generation_service_spec.rb b/spec/enterprise/services/captain/documents/context_generation_service_spec.rb new file mode 100644 index 000000000..248bc19cc --- /dev/null +++ b/spec/enterprise/services/captain/documents/context_generation_service_spec.rb @@ -0,0 +1,108 @@ +require 'rails_helper' +require 'ostruct' + +RSpec.describe Captain::Documents::ContextGenerationService do + let(:chat) { instance_double(RubyLLM::Chat) } + let(:response) { instance_double(RubyLLM::Message, content: " Pricing page context. \n") } + + before do + InstallationConfig.find_or_initialize_by(name: 'CAPTAIN_OPEN_AI_API_KEY').update!(value: 'test-key') + InstallationConfig.where(name: %w[CAPTAIN_CONTEXT_GENERATION_MODEL CAPTAIN_OPEN_AI_MODEL]).delete_all + + allow(RubyLLM).to receive(:chat).and_return(chat) + allow(chat).to receive(:with_temperature).and_return(chat) + allow(chat).to receive(:with_instructions).and_return(chat) + allow(chat).to receive(:ask).and_return(response) + end + + describe '#generate' do + it 'returns stripped context text from the model response' do + service = described_class.new( + document_content: 'Pricing docs for all plans', + chunk_content: 'Business plan starts at $19/agent/month.', + account_id: 1 + ) + + result = service.generate + + expect(result).to eq('Pricing page context.') + end + + it 'uses explicit model when provided' do + service = described_class.new( + document_content: 'Doc text', + chunk_content: 'Chunk text', + account_id: 1, + model: 'gpt-4.1-mini' + ) + + expect(RubyLLM).to receive(:chat).with(model: 'gpt-4.1-mini').and_return(chat) + + service.generate + end + + it 'uses CAPTAIN_CONTEXT_GENERATION_MODEL when configured' do + create(:installation_config, name: 'CAPTAIN_CONTEXT_GENERATION_MODEL', value: 'gpt-4.1-mini') + + service = described_class.new( + document_content: 'Doc text', + chunk_content: 'Chunk text', + account_id: 1 + ) + + expect(RubyLLM).to receive(:chat).with(model: 'gpt-4.1-mini').and_return(chat) + + service.generate + end + + it 'uses CAPTAIN_OPEN_AI_MODEL when context model is not configured' do + create(:installation_config, name: 'CAPTAIN_OPEN_AI_MODEL', value: 'gpt-4o-mini') + + service = described_class.new( + document_content: 'Doc text', + chunk_content: 'Chunk text', + account_id: 1 + ) + + expect(RubyLLM).to receive(:chat).with(model: 'gpt-4o-mini').and_return(chat) + + service.generate + end + + it 'falls back to the service default model when no model config exists' do + service = described_class.new( + document_content: 'Doc text', + chunk_content: 'Chunk text', + account_id: 1 + ) + + expect(RubyLLM).to receive(:chat).with(model: 'gpt-4.1').and_return(chat) + + service.generate + end + + it 'returns empty string when chunk content is blank' do + service = described_class.new( + document_content: 'Doc text', + chunk_content: '', + account_id: 1 + ) + + expect(service.generate).to eq('') + expect(RubyLLM).not_to have_received(:chat) + end + + it 'returns empty string when LLM call fails' do + error = RubyLLM::Error.new(OpenStruct.new(body: {}), 'request failed') + allow(chat).to receive(:ask).and_raise(error) + + service = described_class.new( + document_content: 'Doc text', + chunk_content: 'Chunk text', + account_id: 1 + ) + + expect(service.generate).to eq('') + end + end +end