From 2484a4f07481f8b16cae3de0411f22394e6845db Mon Sep 17 00:00:00 2001 From: Shivam Mishra Date: Wed, 14 Jan 2026 11:32:01 +0530 Subject: [PATCH] feat: use new helpers --- .../conversation/response_builder_job.rb | 3 +-- .../services/captain/copilot/chat_service.rb | 13 +++++++------ .../captain/llm/assistant_chat_service.rb | 9 +++++++-- enterprise/app/services/llm/base_ai_service.rb | 9 ++++++++- .../enterprise/captain/base_task_service.rb | 4 +--- .../conversation/response_builder_job_spec.rb | 3 ++- spec/lib/captain/base_task_service_spec.rb | 18 ------------------ 7 files changed, 26 insertions(+), 33 deletions(-) diff --git a/enterprise/app/jobs/captain/conversation/response_builder_job.rb b/enterprise/app/jobs/captain/conversation/response_builder_job.rb index 0f5e82e7f..b1393a5ba 100644 --- a/enterprise/app/jobs/captain/conversation/response_builder_job.rb +++ b/enterprise/app/jobs/captain/conversation/response_builder_job.rb @@ -47,8 +47,7 @@ class Captain::Conversation::ResponseBuilderJob < ApplicationJob return process_action('handoff') if handoff_requested? create_messages - Rails.logger.info("[CAPTAIN][ResponseBuilderJob] Incrementing response usage for #{account.id}") - account.increment_response_usage + account.increment_response_usage_for_model(account.captain_assistant_model) end def collect_previous_messages diff --git a/enterprise/app/services/captain/copilot/chat_service.rb b/enterprise/app/services/captain/copilot/chat_service.rb index 229bf90ca..196a13c36 100644 --- a/enterprise/app/services/captain/copilot/chat_service.rb +++ b/enterprise/app/services/captain/copilot/chat_service.rb @@ -4,10 +4,10 @@ class Captain::Copilot::ChatService < Llm::BaseAiService attr_reader :assistant, :account, :user, :copilot_thread, :previous_history, :messages def initialize(assistant, config) - super() - @assistant = assistant @account = assistant.account + super(account: @account) + @user = nil @copilot_thread = nil @previous_history = [] @@ -24,16 +24,17 @@ class Captain::Copilot::ChatService < Llm::BaseAiService response = request_chat_completion Rails.logger.debug { "#{self.class.name} Assistant: #{@assistant.id}, Received response #{response}" } - Rails.logger.info( - "#{self.class.name} Assistant: #{@assistant.id}, Incrementing response usage for account #{@account.id}" - ) - @account.increment_response_usage + increment_usage response end private + def setup_model + @model = @account.captain_copilot_model + end + def setup_user(config) @user = @account.users.find_by(id: config[:user_id]) if config[:user_id].present? end diff --git a/enterprise/app/services/captain/llm/assistant_chat_service.rb b/enterprise/app/services/captain/llm/assistant_chat_service.rb index 5a2976e39..ce4770604 100644 --- a/enterprise/app/services/captain/llm/assistant_chat_service.rb +++ b/enterprise/app/services/captain/llm/assistant_chat_service.rb @@ -2,9 +2,10 @@ class Captain::Llm::AssistantChatService < Llm::BaseAiService include Captain::ChatHelper def initialize(assistant: nil, conversation_id: nil) - super() - @assistant = assistant + @account = @assistant&.account + super(account: @account) + @conversation_id = conversation_id @messages = [system_message] @@ -27,6 +28,10 @@ class Captain::Llm::AssistantChatService < Llm::BaseAiService private + def setup_model + @model = @account&.captain_assistant_model || DEFAULT_MODEL + end + def build_tools [Captain::Tools::SearchDocumentationService.new(@assistant, user: nil)] end diff --git a/enterprise/app/services/llm/base_ai_service.rb b/enterprise/app/services/llm/base_ai_service.rb index a5a91cf24..25ec8935b 100644 --- a/enterprise/app/services/llm/base_ai_service.rb +++ b/enterprise/app/services/llm/base_ai_service.rb @@ -8,7 +8,8 @@ class Llm::BaseAiService attr_reader :model, :temperature - def initialize + def initialize(account: nil) + @account = account Llm::Config.initialize! setup_model setup_temperature @@ -18,6 +19,12 @@ class Llm::BaseAiService RubyLLM.chat(model: model).with_temperature(temperature) end + def increment_usage + return unless @account + + @account.increment_response_usage_for_model(@model) + end + private def setup_model diff --git a/enterprise/lib/enterprise/captain/base_task_service.rb b/enterprise/lib/enterprise/captain/base_task_service.rb index c3514758e..76724e7be 100644 --- a/enterprise/lib/enterprise/captain/base_task_service.rb +++ b/enterprise/lib/enterprise/captain/base_task_service.rb @@ -28,8 +28,6 @@ module Enterprise::Captain::BaseTaskService end def increment_usage - credits = account.using_openai_hook_key? ? 1 : Llm::Models.credit_multiplier_for(configured_model) - Rails.logger.info("[CAPTAIN][#{self.class.name}] Incrementing response usage for account #{account.id} by #{credits} credits") - account.increment_response_usage(credits: credits) + account.increment_response_usage_for_model(configured_model) end end diff --git a/spec/enterprise/jobs/captain/conversation/response_builder_job_spec.rb b/spec/enterprise/jobs/captain/conversation/response_builder_job_spec.rb index 777086613..0acbfd08c 100644 --- a/spec/enterprise/jobs/captain/conversation/response_builder_job_spec.rb +++ b/spec/enterprise/jobs/captain/conversation/response_builder_job_spec.rb @@ -1,7 +1,8 @@ require 'rails_helper' RSpec.describe Captain::Conversation::ResponseBuilderJob, type: :job do - let(:account) { create(:account, custom_attributes: { plan_name: 'startups' }) } + # Use gpt-5-mini (credit_multiplier: 1) to keep credit consumption predictable in tests + let(:account) { create(:account, custom_attributes: { plan_name: 'startups' }, captain_models: { 'assistant' => 'gpt-5-mini' }) } let(:inbox) { create(:inbox, account: account) } let(:assistant) { create(:captain_assistant, account: account) } let(:captain_inbox_association) { create(:captain_inbox, captain_assistant: assistant, inbox: inbox) } diff --git a/spec/lib/captain/base_task_service_spec.rb b/spec/lib/captain/base_task_service_spec.rb index 1007796b2..661bf7f0b 100644 --- a/spec/lib/captain/base_task_service_spec.rb +++ b/spec/lib/captain/base_task_service_spec.rb @@ -254,24 +254,6 @@ RSpec.describe Captain::BaseTaskService do end end - describe '#api_key' do - context 'when openai hook is configured' do - let(:hook) { create(:integrations_hook, account: account, app_id: 'openai', status: 'enabled', settings: { 'api_key' => 'hook-key' }) } - - before { hook } - - it 'uses api key from hook' do - expect(service.send(:api_key)).to eq('hook-key') - end - end - - context 'when openai hook is not configured' do - it 'uses system api key' do - expect(service.send(:api_key)).to eq('test-key') - end - end - end - describe '#prompt_from_file' do it 'reads prompt from file' do allow(Rails.root).to receive(:join).and_return(instance_double(Pathname, read: 'Test prompt content'))