-
-
Date: Thu, 23 Jan 2025 01:23:18 +0530
Subject: [PATCH 4/4] feat: setup captain limits (#10713)
This pull request introduces several changes to implement and manage
usage limits for the Captain AI service. The key changes include adding
configuration for plan limits, updating error messages, modifying
controllers and models to handle usage limits, and updating tests to
ensure the new functionality works correctly.
## Implementation Checklist
- [x] Ability to configure captain limits per check
- [x] Update response for `usage_limits` to include captain limits
- [x] Methods to increment or reset captain responses limits in the
`limits` column for the `Account` model
- [x] Check documents limit using a count query
- [x] Ensure Captain hand-off if a limit is reached
- [x] Ensure limits are enforced for Copilot Chat
- [x] Ensure limits are reset when stripe webhook comes in
- [x] Increment usage for FAQ generation and Contact notes
- [x] Ensure documents limit is enforced
These changes ensure that the Captain AI service operates within the defined usage limits for different subscription plans, providing appropriate error messages and handling when limits are exceeded.
---
app/fields/enterprise/account_limits_field.rb | 2 +-
config/installation_config.yml | 6 +
config/locales/en.yml | 3 +-
.../accounts/captain/documents_controller.rb | 2 +
.../v1/accounts/conversations_controller.rb | 1 +
.../super_admin/app_configs_controller.rb | 3 +-
enterprise/app/helpers/captain/chat_helper.rb | 1 -
.../conversation/response_builder_job.rb | 3 +
.../tools/simple_page_crawl_parser_job.rb | 14 ++
enterprise/app/models/captain/document.rb | 13 ++
enterprise/app/models/enterprise/account.rb | 82 +++++++++-
enterprise/app/models/enterprise/inbox.rb | 10 +-
.../services/captain/copilot/chat_service.rb | 6 +-
.../billing/handle_stripe_event_service.rb | 5 +
.../hook_execution_service.rb | 14 ++
enterprise/listeners/captain_listener.rb | 3 +-
.../captain/documents_controller_spec.rb | 22 ++-
.../conversation/response_builder_job_spec.rb | 34 ++++
spec/enterprise/models/account_spec.rb | 147 ++++++++++++++----
.../captain/copilot/chat_service_spec.rb | 34 ++++
.../handle_stripe_event_service_spec.rb | 43 +++--
spec/models/account_spec.rb | 3 +-
22 files changed, 394 insertions(+), 57 deletions(-)
create mode 100644 spec/enterprise/jobs/captain/conversation/response_builder_job_spec.rb
create mode 100644 spec/enterprise/services/captain/copilot/chat_service_spec.rb
diff --git a/app/fields/enterprise/account_limits_field.rb b/app/fields/enterprise/account_limits_field.rb
index 5833952ad..b014435c6 100644
--- a/app/fields/enterprise/account_limits_field.rb
+++ b/app/fields/enterprise/account_limits_field.rb
@@ -2,6 +2,6 @@ require 'administrate/field/base'
class Enterprise::AccountLimitsField < Administrate::Field::Base
def to_s
- data.present? ? data.to_json : { agents: nil, inboxes: nil }.to_json
+ data.present? ? data.to_json : { agents: nil, inboxes: nil, captain_responses: nil, captain_documents: nil }.to_json
end
end
diff --git a/config/installation_config.yml b/config/installation_config.yml
index 2b4241697..a3b15a191 100644
--- a/config/installation_config.yml
+++ b/config/installation_config.yml
@@ -142,6 +142,12 @@
display_title: 'OpenAI Model'
description: 'The OpenAI model configured for use in Captain AI. Default: gpt-4o-mini'
locked: false
+- name: CAPTAIN_CLOUD_PLAN_LIMITS
+ display_title: 'Captain Cloud Plan Limits'
+ description: 'The limits for the Captain AI service for different plans'
+ value:
+ type: code
+
# End of Captain Config
# ------- Chatwoot Internal Config for Cloud ----#
diff --git a/config/locales/en.yml b/config/locales/en.yml
index 310235d5e..c525074cf 100644
--- a/config/locales/en.yml
+++ b/config/locales/en.yml
@@ -230,7 +230,8 @@ en:
name: 'Linear'
description: 'Create issues in Linear directly from your conversation window. Alternatively, link existing Linear issues for a more streamlined and efficient issue tracking process.'
captain:
- copilot_error: 'Please connect an assistant to this inbox to use copilot'
+ copilot_error: 'Please connect an assistant to this inbox to use Copilot'
+ copilot_limit: 'You are out of Copilot credits. You can buy more credits from the billing section.'
public_portal:
search:
search_placeholder: Search for article by title or body...
diff --git a/enterprise/app/controllers/api/v1/accounts/captain/documents_controller.rb b/enterprise/app/controllers/api/v1/accounts/captain/documents_controller.rb
index 9743c3892..594aa0642 100644
--- a/enterprise/app/controllers/api/v1/accounts/captain/documents_controller.rb
+++ b/enterprise/app/controllers/api/v1/accounts/captain/documents_controller.rb
@@ -23,6 +23,8 @@ class Api::V1::Accounts::Captain::DocumentsController < Api::V1::Accounts::BaseC
@document = @assistant.documents.build(document_params)
@document.save!
+ rescue Captain::Document::LimitExceededError => e
+ render_could_not_create_error(e.message)
end
def destroy
diff --git a/enterprise/app/controllers/enterprise/api/v1/accounts/conversations_controller.rb b/enterprise/app/controllers/enterprise/api/v1/accounts/conversations_controller.rb
index f22812f19..e9113f4d0 100644
--- a/enterprise/app/controllers/enterprise/api/v1/accounts/conversations_controller.rb
+++ b/enterprise/app/controllers/enterprise/api/v1/accounts/conversations_controller.rb
@@ -7,6 +7,7 @@ module Enterprise::Api::V1::Accounts::ConversationsController
def copilot
assistant = @conversation.inbox.captain_assistant
return render json: { message: I18n.t('captain.copilot_error') } unless assistant
+ return render json: { message: I18n.t('captain.copilot_limit') } unless @conversation.inbox.captain_active?
response = Captain::Copilot::ChatService.new(
assistant,
diff --git a/enterprise/app/controllers/enterprise/super_admin/app_configs_controller.rb b/enterprise/app/controllers/enterprise/super_admin/app_configs_controller.rb
index 8cf8e7258..b11aab487 100644
--- a/enterprise/app/controllers/enterprise/super_admin/app_configs_controller.rb
+++ b/enterprise/app/controllers/enterprise/super_admin/app_configs_controller.rb
@@ -32,6 +32,7 @@ module Enterprise::SuperAdmin::AppConfigsController
end
def internal_config_options
- %w[CHATWOOT_INBOX_TOKEN CHATWOOT_INBOX_HMAC_KEY ANALYTICS_TOKEN CLEARBIT_API_KEY DASHBOARD_SCRIPTS BLOCKED_EMAIL_DOMAINS]
+ %w[CHATWOOT_INBOX_TOKEN CHATWOOT_INBOX_HMAC_KEY ANALYTICS_TOKEN CLEARBIT_API_KEY DASHBOARD_SCRIPTS BLOCKED_EMAIL_DOMAINS
+ CAPTAIN_CLOUD_PLAN_LIMITS]
end
end
diff --git a/enterprise/app/helpers/captain/chat_helper.rb b/enterprise/app/helpers/captain/chat_helper.rb
index 258328f9d..a10c9e797 100644
--- a/enterprise/app/helpers/captain/chat_helper.rb
+++ b/enterprise/app/helpers/captain/chat_helper.rb
@@ -35,7 +35,6 @@ module Captain::ChatHelper
def handle_response(response)
message = response.dig('choices', 0, 'message')
-
if message['tool_calls']
process_tool_calls(message['tool_calls'])
else
diff --git a/enterprise/app/jobs/captain/conversation/response_builder_job.rb b/enterprise/app/jobs/captain/conversation/response_builder_job.rb
index 6b6570f0b..c8c511a57 100644
--- a/enterprise/app/jobs/captain/conversation/response_builder_job.rb
+++ b/enterprise/app/jobs/captain/conversation/response_builder_job.rb
@@ -3,6 +3,7 @@ class Captain::Conversation::ResponseBuilderJob < ApplicationJob
def perform(conversation, assistant)
@conversation = conversation
+ @inbox = conversation.inbox
@assistant = assistant
ActiveRecord::Base.transaction do
@@ -25,6 +26,8 @@ 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
end
def collect_previous_messages
diff --git a/enterprise/app/jobs/captain/tools/simple_page_crawl_parser_job.rb b/enterprise/app/jobs/captain/tools/simple_page_crawl_parser_job.rb
index fa658dd80..ccf34adf2 100644
--- a/enterprise/app/jobs/captain/tools/simple_page_crawl_parser_job.rb
+++ b/enterprise/app/jobs/captain/tools/simple_page_crawl_parser_job.rb
@@ -3,6 +3,13 @@ class Captain::Tools::SimplePageCrawlParserJob < ApplicationJob
def perform(assistant_id:, page_link:)
assistant = Captain::Assistant.find(assistant_id)
+ account = assistant.account
+
+ if limit_exceeded?(account)
+ Rails.logger.info("Document limit exceeded for #{assistant_id}")
+ return
+ end
+
crawler = Captain::Tools::SimplePageCrawlService.new(page_link)
page_title = crawler.page_title || ''
@@ -18,4 +25,11 @@ class Captain::Tools::SimplePageCrawlParserJob < ApplicationJob
rescue StandardError => e
raise "Failed to parse data: #{page_link} #{e.message}"
end
+
+ private
+
+ def limit_exceeded?(account)
+ limits = account.usage_limits[:captain][:documents]
+ limits[:current_available].negative? || limits[:current_available].zero?
+ end
end
diff --git a/enterprise/app/models/captain/document.rb b/enterprise/app/models/captain/document.rb
index 11338d409..8331d291b 100644
--- a/enterprise/app/models/captain/document.rb
+++ b/enterprise/app/models/captain/document.rb
@@ -20,6 +20,7 @@
# index_captain_documents_on_status (status)
#
class Captain::Document < ApplicationRecord
+ class LimitExceededError < StandardError; end
self.table_name = 'captain_documents'
belongs_to :assistant, class_name: 'Captain::Assistant'
@@ -35,7 +36,10 @@ class Captain::Document < ApplicationRecord
available: 1
}
+ 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
scope :ordered, -> { order(created_at: :desc) }
@@ -56,7 +60,16 @@ class Captain::Document < ApplicationRecord
Captain::Documents::ResponseBuilderJob.perform_later(self)
end
+ def update_document_usage
+ account.update_document_usage
+ end
+
def ensure_account_id
self.account_id = assistant&.account_id
end
+
+ def ensure_within_plan_limit
+ limits = account.usage_limits[:captain][:documents]
+ raise LimitExceededError, 'Document limit exceeded' unless limits[:current_available].positive?
+ end
end
diff --git a/enterprise/app/models/enterprise/account.rb b/enterprise/app/models/enterprise/account.rb
index bc07e1171..a1b0d0449 100644
--- a/enterprise/app/models/enterprise/account.rb
+++ b/enterprise/app/models/enterprise/account.rb
@@ -1,11 +1,37 @@
module Enterprise::Account
+ CAPTAIN_RESPONSES = 'captain_responses'.freeze
+ CAPTAIN_DOCUMENTS = 'captain_documents'.freeze
+ CAPTAIN_RESPONSES_USAGE = 'captain_responses_usage'.freeze
+ CAPTAIN_DOCUMENTS_USAGE = 'captain_documents_usage'.freeze
+
def usage_limits
{
agents: agent_limits.to_i,
- inboxes: get_limits(:inboxes).to_i
+ inboxes: get_limits(:inboxes).to_i,
+ captain: {
+ documents: get_captain_limits(:documents),
+ responses: get_captain_limits(:responses)
+ }
}
end
+ def increment_response_usage
+ current_usage = custom_attributes[CAPTAIN_RESPONSES_USAGE].to_i || 0
+ custom_attributes[CAPTAIN_RESPONSES_USAGE] = current_usage + 1
+ save
+ end
+
+ def reset_response_usage
+ custom_attributes[CAPTAIN_RESPONSES_USAGE] = 0
+ save
+ end
+
+ def update_document_usage
+ # this will ensure that the document count is always accurate
+ custom_attributes[CAPTAIN_DOCUMENTS_USAGE] = captain_documents.count
+ save
+ end
+
def subscribed_features
plan_features = InstallationConfig.find_by(name: 'CHATWOOT_CLOUD_PLAN_FEATURES')&.value
return [] if plan_features.blank?
@@ -13,8 +39,58 @@ module Enterprise::Account
plan_features[plan_name]
end
+ def captain_monthly_limit
+ default_limits = default_captain_limits
+
+ {
+ documents: self[:limits][CAPTAIN_DOCUMENTS] || default_limits['documents'],
+ responses: self[:limits][CAPTAIN_RESPONSES] || default_limits['responses']
+ }.with_indifferent_access
+ end
+
private
+ def get_captain_limits(type)
+ total_count = captain_monthly_limit[type.to_s].to_i
+
+ consumed = if type == :documents
+ custom_attributes[CAPTAIN_DOCUMENTS_USAGE].to_i || 0
+ else
+ custom_attributes[CAPTAIN_RESPONSES_USAGE].to_i || 0
+ end
+
+ consumed = 0 if consumed.negative?
+
+ {
+ total_count: total_count,
+ current_available: (total_count - consumed).clamp(0, total_count),
+ consumed: consumed
+ }
+ end
+
+ def default_captain_limits
+ max_limits = { documents: ChatwootApp.max_limit, responses: ChatwootApp.max_limit }.with_indifferent_access
+ zero_limits = { documents: 0, responses: 0 }.with_indifferent_access
+ plan_quota = InstallationConfig.find_by(name: 'CAPTAIN_CLOUD_PLAN_LIMITS')&.value
+
+ # If there are no limits configured, we allow max usage
+ return max_limits if plan_quota.blank?
+
+ # if there is plan_quota configred, but plan_name is not present, we return zero limits
+ return zero_limits if plan_name.blank?
+
+ begin
+ # Now we parse the plan_quota and return the limits for the plan name
+ # but if there's no plan_name present in the plan_quota, we return zero limits
+ plan_quota = JSON.parse(plan_quota) if plan_quota.present?
+ plan_quota[plan_name.downcase] || zero_limits
+ rescue StandardError
+ # if there's any error in parsing the plan_quota, we return max limits
+ # this is to ensure that we don't block the user from using the product
+ max_limits
+ end
+ end
+
def plan_name
custom_attributes['plan_name']
end
@@ -41,7 +117,9 @@ module Enterprise::Account
'type' => 'object',
'properties' => {
'inboxes' => { 'type': 'number' },
- 'agents' => { 'type': 'number' }
+ 'agents' => { 'type': 'number' },
+ 'captain_responses' => { 'type': 'number' },
+ 'captain_documents' => { 'type': 'number' }
},
'required' => [],
'additionalProperties' => false
diff --git a/enterprise/app/models/enterprise/inbox.rb b/enterprise/app/models/enterprise/inbox.rb
index 5f89a92a8..4cd627e41 100644
--- a/enterprise/app/models/enterprise/inbox.rb
+++ b/enterprise/app/models/enterprise/inbox.rb
@@ -6,11 +6,19 @@ module Enterprise::Inbox
end
def active_bot?
- super || captain_assistant.present?
+ super || captain_active?
+ end
+
+ def captain_active?
+ captain_assistant.present? && more_responses?
end
private
+ def more_responses?
+ account.usage_limits[:captain][:responses][:current_available].positive?
+ end
+
def get_agent_ids_over_assignment_limit(limit)
conversations.open.select(:assignee_id).group(:assignee_id).having("count(*) >= #{limit.to_i}").filter_map(&:assignee_id)
end
diff --git a/enterprise/app/services/captain/copilot/chat_service.rb b/enterprise/app/services/captain/copilot/chat_service.rb
index 6bcda66fa..019136540 100644
--- a/enterprise/app/services/captain/copilot/chat_service.rb
+++ b/enterprise/app/services/captain/copilot/chat_service.rb
@@ -15,7 +15,11 @@ class Captain::Copilot::ChatService < Captain::Llm::BaseOpenAiService
def generate_response(input)
@messages << { role: 'user', content: input } if input.present?
- request_chat_completion
+ response = request_chat_completion
+ Rails.logger.info("[CAPTAIN][CopilotChatService] Incrementing response usage for #{@assistant.account.id}")
+ @assistant.account.increment_response_usage
+
+ response
end
private
diff --git a/enterprise/app/services/enterprise/billing/handle_stripe_event_service.rb b/enterprise/app/services/enterprise/billing/handle_stripe_event_service.rb
index 1ddf00d7a..c3c570596 100644
--- a/enterprise/app/services/enterprise/billing/handle_stripe_event_service.rb
+++ b/enterprise/app/services/enterprise/billing/handle_stripe_event_service.rb
@@ -22,6 +22,7 @@ class Enterprise::Billing::HandleStripeEventService
update_account_attributes(subscription, plan)
change_plan_features
+ reset_captain_usage
end
def update_account_attributes(subscription, plan)
@@ -56,6 +57,10 @@ class Enterprise::Billing::HandleStripeEventService
account.save!
end
+ def reset_captain_usage
+ account.reset_response_usage
+ end
+
def ensure_event_context(event)
@event = event
end
diff --git a/enterprise/app/services/enterprise/message_templates/hook_execution_service.rb b/enterprise/app/services/enterprise/message_templates/hook_execution_service.rb
index 858b5e903..92ccba553 100644
--- a/enterprise/app/services/enterprise/message_templates/hook_execution_service.rb
+++ b/enterprise/app/services/enterprise/message_templates/hook_execution_service.rb
@@ -2,6 +2,7 @@ module Enterprise::MessageTemplates::HookExecutionService
def trigger_templates
super
return unless should_process_captain_response?
+ return perform_handoff unless inbox.captain_active?
Captain::Conversation::ResponseBuilderJob.perform_later(
conversation,
@@ -12,4 +13,17 @@ module Enterprise::MessageTemplates::HookExecutionService
def should_process_captain_response?
conversation.pending? && message.incoming? && inbox.captain_assistant.present?
end
+
+ def perform_handoff
+ return unless conversation.pending?
+
+ Rails.logger.info("Captain limit exceeded, performing handoff mid-conversation for conversation: #{conversation.id}")
+ conversation.messages.create!(
+ message_type: :outgoing,
+ account_id: conversation.account.id,
+ inbox_id: conversation.inbox.id,
+ content: 'Transferring to another agent for further assistance.'
+ )
+ conversation.bot_handoff!
+ end
end
diff --git a/enterprise/listeners/captain_listener.rb b/enterprise/listeners/captain_listener.rb
index da77bfbaa..adf156562 100644
--- a/enterprise/listeners/captain_listener.rb
+++ b/enterprise/listeners/captain_listener.rb
@@ -2,8 +2,7 @@ class CaptainListener < BaseListener
def conversation_resolved(event)
conversation = extract_conversation_and_account(event)[0]
assistant = conversation.inbox.captain_assistant
-
- return if assistant.blank?
+ return unless conversation.inbox.captain_active?
Captain::Llm::ContactNotesService.new(assistant, conversation).generate_and_update_notes if assistant.config['feature_memory'].present?
Captain::Llm::ConversationFaqService.new(assistant, conversation).generate_and_deduplicate if assistant.config['feature_faq'].present?
diff --git a/spec/enterprise/controllers/api/v1/accounts/captain/documents_controller_spec.rb b/spec/enterprise/controllers/api/v1/accounts/captain/documents_controller_spec.rb
index f9385af1c..ef0aa6b2e 100644
--- a/spec/enterprise/controllers/api/v1/accounts/captain/documents_controller_spec.rb
+++ b/spec/enterprise/controllers/api/v1/accounts/captain/documents_controller_spec.rb
@@ -1,12 +1,17 @@
require 'rails_helper'
RSpec.describe 'Api::V1::Accounts::Captain::Documents', type: :request do
- let(:account) { create(:account) }
+ let(:account) { create(:account, custom_attributes: { plan_name: 'startups' }) }
let(:admin) { create(:user, account: account, role: :administrator) }
let(:agent) { create(:user, account: account, role: :agent) }
let(:assistant) { create(:captain_assistant, account: account) }
let(:assistant2) { create(:captain_assistant, account: account) }
let(:document) { create(:captain_document, assistant: assistant, account: account) }
+ let(:captain_limits) do
+ {
+ :startups => { :documents => 1, :responses => 100 }
+ }.with_indifferent_access
+ end
def json_response
JSON.parse(response.body, symbolize_names: true)
@@ -212,6 +217,21 @@ RSpec.describe 'Api::V1::Accounts::Captain::Documents', type: :request do
expect(response).to have_http_status(:unprocessable_entity)
end
end
+
+ context 'with limits exceeded' do
+ before do
+ create_list(:captain_document, 5, assistant: assistant, account: account)
+
+ create(:installation_config, name: 'CAPTAIN_CLOUD_PLAN_LIMITS', value: captain_limits.to_json)
+ post "/api/v1/accounts/#{account.id}/captain/documents",
+ params: valid_attributes,
+ headers: admin.create_new_auth_token
+ end
+
+ it 'returns an error' do
+ expect(response).to have_http_status(:unprocessable_entity)
+ end
+ end
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
new file mode 100644
index 000000000..1e4a6e824
--- /dev/null
+++ b/spec/enterprise/jobs/captain/conversation/response_builder_job_spec.rb
@@ -0,0 +1,34 @@
+require 'rails_helper'
+
+RSpec.describe Captain::Conversation::ResponseBuilderJob, type: :job do
+ let(:account) { create(:account, custom_attributes: { plan_name: 'startups' }) }
+ 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) }
+
+ describe '#perform' do
+ let(:conversation) { create(:conversation, inbox: inbox, account: account) }
+ let(:mock_llm_chat_service) { instance_double(Captain::Llm::AssistantChatService) }
+
+ before do
+ create(:message, conversation: conversation, content: 'Hello', message_type: :incoming)
+
+ allow(inbox).to receive(:captain_active?).and_return(true)
+ allow(Captain::Llm::AssistantChatService).to receive(:new).and_return(mock_llm_chat_service)
+ allow(mock_llm_chat_service).to receive(:generate_response).and_return({ 'response' => 'Hey, welcome to Captain Specs' })
+ end
+
+ it 'generates and processes response' do
+ described_class.perform_now(conversation, assistant)
+ expect(conversation.messages.count).to eq(2)
+ expect(conversation.messages.outgoing.count).to eq(1)
+ expect(conversation.messages.last.content).to eq('Hey, welcome to Captain Specs')
+ end
+
+ it 'increments usage response' do
+ described_class.perform_now(conversation, assistant)
+ account.reload
+ expect(account.usage_limits[:captain][:responses][:consumed]).to eq(1)
+ end
+ end
+end
diff --git a/spec/enterprise/models/account_spec.rb b/spec/enterprise/models/account_spec.rb
index 8e2681a9b..e36a7217d 100644
--- a/spec/enterprise/models/account_spec.rb
+++ b/spec/enterprise/models/account_spec.rb
@@ -27,12 +27,120 @@ RSpec.describe Account, type: :model do
end
end
- describe 'usage_limits' do
+ context 'with usage_limits' do
+ let(:captain_limits) do
+ {
+ :startups => { :documents => 100, :responses => 100 },
+ :business => { :documents => 200, :responses => 300 },
+ :enterprise => { :documents => 300, :responses => 500 }
+ }.with_indifferent_access
+ end
+ let(:account) { create(:account, { custom_attributes: { plan_name: 'startups' } }) }
+ let(:assistant) { create(:captain_assistant, account: account) }
+
before do
create(:installation_config, name: 'ACCOUNT_AGENTS_LIMIT', value: 20)
end
- let!(:account) { create(:account) }
+ describe 'when captain limits are configured' do
+ before do
+ create_list(:captain_document, 3, account: account, assistant: assistant, status: :available)
+ create(:installation_config, name: 'CAPTAIN_CLOUD_PLAN_LIMITS', value: captain_limits.to_json)
+ end
+
+ ## Document
+ it 'updates document count accurately' do
+ account.update_document_usage
+ expect(account.custom_attributes['captain_documents_usage']).to eq(3)
+ end
+
+ it 'handles zero documents' do
+ account.captain_documents.destroy_all
+ account.update_document_usage
+ expect(account.custom_attributes['captain_documents_usage']).to eq(0)
+ end
+
+ it 'reflects document limits' do
+ document_limits = account.usage_limits[:captain][:documents]
+
+ expect(document_limits[:consumed]).to eq 3
+ expect(document_limits[:current_available]).to eq captain_limits[:startups][:documents] - 3
+ end
+
+ ## Responses
+ it 'incrementing responses updates usage_limits' do
+ account.increment_response_usage
+
+ responses_limits = account.usage_limits[:captain][:responses]
+
+ expect(account.custom_attributes['captain_responses_usage']).to eq 1
+ expect(responses_limits[:consumed]).to eq 1
+ expect(responses_limits[:current_available]).to eq captain_limits[:startups][:responses] - 1
+ end
+
+ it 'reseting responses limits updates usage_limits' do
+ account.custom_attributes['captain_responses_usage'] = 30
+ account.save!
+
+ responses_limits = account.usage_limits[:captain][:responses]
+
+ expect(responses_limits[:consumed]).to eq 30
+ expect(responses_limits[:current_available]).to eq captain_limits[:startups][:responses] - 30
+
+ account.reset_response_usage
+ responses_limits = account.usage_limits[:captain][:responses]
+
+ expect(account.custom_attributes['captain_responses_usage']).to eq 0
+ expect(responses_limits[:consumed]).to eq 0
+ expect(responses_limits[:current_available]).to eq captain_limits[:startups][:responses]
+ end
+
+ it 'returns monthly limit accurately' do
+ %w[startups business enterprise].each do |plan|
+ account.custom_attributes = { 'plan_name': plan }
+ account.save!
+ expect(account.captain_monthly_limit).to eq captain_limits[plan]
+ end
+ end
+
+ it 'current_available is never out of bounds' do
+ account.custom_attributes['captain_responses_usage'] = 3000
+ account.save!
+
+ responses_limits = account.usage_limits[:captain][:responses]
+ expect(responses_limits[:consumed]).to eq 3000
+ expect(responses_limits[:current_available]).to eq 0
+
+ account.custom_attributes['captain_responses_usage'] = -100
+ account.save!
+
+ responses_limits = account.usage_limits[:captain][:responses]
+ expect(responses_limits[:consumed]).to eq 0
+ expect(responses_limits[:current_available]).to eq captain_limits[:startups][:responses]
+ end
+ end
+
+ describe 'when captain limits are not configured' do
+ it 'returns default values' do
+ account.custom_attributes = { 'plan_name': 'unknown' }
+ expect(account.captain_monthly_limit).to eq(
+ { documents: ChatwootApp.max_limit, responses: ChatwootApp.max_limit }.with_indifferent_access
+ )
+ end
+ end
+
+ describe 'when limits are configured for an account' do
+ before do
+ create(:installation_config, name: 'CAPTAIN_CLOUD_PLAN_LIMITS', value: captain_limits.to_json)
+ account.update(limits: { captain_documents: 5555, captain_responses: 9999 })
+ end
+
+ it 'returns limits based on custom attributes' do
+ usage_limits = account.usage_limits
+ expect(usage_limits[:captain][:documents][:total_count]).to eq(5555)
+ expect(usage_limits[:captain][:responses][:total_count]).to eq(9999)
+ end
+ end
describe 'audit logs' do
it 'returns audit logs' do
@@ -47,54 +155,29 @@ RSpec.describe Account, type: :model do
end
it 'returns max limits from global config when enterprise version' do
- expect(account.usage_limits).to eq(
- {
- agents: 20,
- inboxes: ChatwootApp.max_limit
- }
- )
+ expect(account.usage_limits[:agents]).to eq(20)
end
it 'returns max limits from account when enterprise version' do
account.update(limits: { agents: 10 })
- expect(account.usage_limits).to eq(
- {
- agents: 10,
- inboxes: ChatwootApp.max_limit
- }
- )
+ expect(account.usage_limits[:agents]).to eq(10)
end
it 'returns limits based on subscription' do
account.update(limits: { agents: 10 }, custom_attributes: { subscribed_quantity: 5 })
- expect(account.usage_limits).to eq(
- {
- agents: 5,
- inboxes: ChatwootApp.max_limit
- }
- )
+ expect(account.usage_limits[:agents]).to eq(5)
end
it 'returns max limits from global config if account limit is absent' do
account.update(limits: { agents: '' })
- expect(account.usage_limits).to eq(
- {
- agents: 20,
- inboxes: ChatwootApp.max_limit
- }
- )
+ expect(account.usage_limits[:agents]).to eq(20)
end
it 'returns max limits from app limit if account limit and installation config is absent' do
account.update(limits: { agents: '' })
InstallationConfig.where(name: 'ACCOUNT_AGENTS_LIMIT').update(value: '')
- expect(account.usage_limits).to eq(
- {
- agents: ChatwootApp.max_limit,
- inboxes: ChatwootApp.max_limit
- }
- )
+ expect(account.usage_limits[:agents]).to eq(ChatwootApp.max_limit)
end
end
diff --git a/spec/enterprise/services/captain/copilot/chat_service_spec.rb b/spec/enterprise/services/captain/copilot/chat_service_spec.rb
new file mode 100644
index 000000000..c4b759c85
--- /dev/null
+++ b/spec/enterprise/services/captain/copilot/chat_service_spec.rb
@@ -0,0 +1,34 @@
+require 'rails_helper'
+
+RSpec.describe Captain::Copilot::ChatService do
+ let(:account) { create(:account, custom_attributes: { plan_name: 'startups' }) }
+ 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) }
+
+ let(:mock_captain_agent) { instance_double(Captain::Agent) }
+ let(:mock_captain_tool) { instance_double(Captain::Tool) }
+ let(:mock_openai_client) { instance_double(OpenAI::Client) }
+
+ describe '#execute' do
+ before do
+ create(:installation_config) { create(:installation_config, name: 'CAPTAIN_OPEN_AI_API_KEY', value: 'test-key') }
+ allow(OpenAI::Client).to receive(:new).and_return(mock_openai_client)
+ allow(mock_openai_client).to receive(:chat).and_return({ choices: [{ message: { content: '{ "result": "Hey" }' } }] }.with_indifferent_access)
+
+ allow(Captain::Agent).to receive(:new).and_return(mock_captain_agent)
+ allow(mock_captain_agent).to receive(:execute).and_return(true)
+ allow(mock_captain_agent).to receive(:register_tool).and_return(true)
+
+ allow(Captain::Tool).to receive(:new).and_return(mock_captain_tool)
+ allow(mock_captain_tool).to receive(:register_method).and_return(true)
+
+ allow(account).to receive(:increment_response_usage).and_return(true)
+ end
+
+ it 'increments usage' do
+ described_class.new(assistant, { previous_messages: ['Hello'], conversation_history: 'Hi' }).generate_response('Hey')
+ expect(account).to have_received(:increment_response_usage).once
+ end
+ end
+end
diff --git a/spec/enterprise/services/enterprise/billing/handle_stripe_event_service_spec.rb b/spec/enterprise/services/enterprise/billing/handle_stripe_event_service_spec.rb
index 36c163f79..d6d2e187e 100644
--- a/spec/enterprise/services/enterprise/billing/handle_stripe_event_service_spec.rb
+++ b/spec/enterprise/services/enterprise/billing/handle_stripe_event_service_spec.rb
@@ -37,19 +37,34 @@ describe Enterprise::Billing::HandleStripeEventService do
end
describe '#perform' do
- it 'handle customer.subscription.updated' do
- allow(event).to receive(:type).and_return('customer.subscription.updated')
- allow(subscription).to receive(:customer).and_return('cus_123')
- stripe_event_service.new.perform(event: event)
- expect(account.reload.custom_attributes).to eq({
- 'stripe_customer_id' => 'cus_123',
- 'stripe_price_id' => 'test',
- 'stripe_product_id' => 'plan_id',
- 'plan_name' => 'Hacker',
- 'subscribed_quantity' => '10',
- 'subscription_ends_on' => Time.zone.at(1_686_567_520).as_json,
- 'subscription_status' => 'active'
- })
+ context 'when it gets customer.subscription.updated event' do
+ it 'updates subscription attributes' do
+ allow(event).to receive(:type).and_return('customer.subscription.updated')
+ allow(subscription).to receive(:customer).and_return('cus_123')
+ stripe_event_service.new.perform(event: event)
+
+ expect(account.reload.custom_attributes).to eq({
+ 'captain_responses_usage' => 0,
+ 'stripe_customer_id' => 'cus_123',
+ 'stripe_price_id' => 'test',
+ 'stripe_product_id' => 'plan_id',
+ 'plan_name' => 'Hacker',
+ 'subscribed_quantity' => '10',
+ 'subscription_ends_on' => Time.zone.at(1_686_567_520).as_json,
+ 'subscription_status' => 'active'
+ })
+ end
+
+ it 'resets captain usage' do
+ 5.times { account.increment_response_usage }
+ expect(account.custom_attributes['captain_responses_usage']).to eq(5)
+
+ allow(event).to receive(:type).and_return('customer.subscription.updated')
+ allow(subscription).to receive(:customer).and_return('cus_123')
+ stripe_event_service.new.perform(event: event)
+
+ expect(account.reload.custom_attributes['captain_responses_usage']).to eq(0)
+ end
end
it 'disable features on customer.subscription.updated for default plan' do
@@ -57,6 +72,7 @@ describe Enterprise::Billing::HandleStripeEventService do
allow(subscription).to receive(:customer).and_return('cus_123')
stripe_event_service.new.perform(event: event)
expect(account.reload.custom_attributes).to eq({
+ 'captain_responses_usage' => 0,
'stripe_customer_id' => 'cus_123',
'stripe_price_id' => 'test',
'stripe_product_id' => 'plan_id',
@@ -96,6 +112,7 @@ describe Enterprise::Billing::HandleStripeEventService do
allow(subscription).to receive(:customer).and_return('cus_123')
stripe_event_service.new.perform(event: event)
expect(account.reload.custom_attributes).to eq({
+ 'captain_responses_usage' => 0,
'stripe_customer_id' => 'cus_123',
'stripe_price_id' => 'test',
'stripe_product_id' => 'plan_id_2',
diff --git a/spec/models/account_spec.rb b/spec/models/account_spec.rb
index aa0556492..4b200fb9a 100644
--- a/spec/models/account_spec.rb
+++ b/spec/models/account_spec.rb
@@ -43,7 +43,8 @@ RSpec.describe Account do
let(:account) { create(:account) }
it 'returns ChatwootApp.max limits' do
- expect(account.usage_limits).to eq({ agents: ChatwootApp.max_limit, inboxes: ChatwootApp.max_limit })
+ expect(account.usage_limits[:agents]).to eq(ChatwootApp.max_limit)
+ expect(account.usage_limits[:inboxes]).to eq(ChatwootApp.max_limit)
end
end